

<!DOCTYPE html>
<html class="writer-html5" lang="en" >
<head>
  <meta charset="utf-8" />
  
  <meta name="viewport" content="width=device-width, initial-scale=1.0" />
  
  <title>mindspore.ops.composite.math_ops &mdash; MindSpore master documentation</title>
  

  
  <link rel="stylesheet" href="../../../../_static/css/theme.css" type="text/css" />
  <link rel="stylesheet" href="../../../../_static/pygments.css" type="text/css" />

  
  

  
  

  

  
  <!--[if lt IE 9]>
    <script src="../../../../_static/js/html5shiv.min.js"></script>
  <![endif]-->
  
    
      <script type="text/javascript" id="documentation_options" data-url_root="../../../../" src="../../../../_static/documentation_options.js"></script>
        <script src="../../../../_static/jquery.js"></script>
        <script src="../../../../_static/underscore.js"></script>
        <script src="../../../../_static/doctools.js"></script>
        <script src="../../../../_static/language_data.js"></script>
        <script async="async" src="https://cdnjs.cloudflare.com/ajax/libs/mathjax/2.7.5/latest.js?config=TeX-AMS-MML_HTMLorMML"></script>
    
    <script type="text/javascript" src="../../../../_static/js/theme.js"></script>

    
    <link rel="index" title="Index" href="../../../../genindex.html" />
    <link rel="search" title="Search" href="../../../../search.html" /> 
</head>

<body class="wy-body-for-nav">

   
  <div class="wy-grid-for-nav">
    
    <nav data-toggle="wy-nav-shift" class="wy-nav-side">
      <div class="wy-side-scroll">
        <div class="wy-side-nav-search" >
          

          
            <a href="../../../../index.html" class="icon icon-home"> MindSpore
          

          
          </a>

          
            
            
          

          
<div role="search">
  <form id="rtd-search-form" class="wy-form" action="../../../../search.html" method="get">
    <input type="text" name="q" placeholder="Search docs" />
    <input type="hidden" name="check_keywords" value="yes" />
    <input type="hidden" name="area" value="default" />
  </form>
</div>

          
        </div>

        
        <div class="wy-menu wy-menu-vertical" data-spy="affix" role="navigation" aria-label="main navigation">
          
            
            
              
            
            
              <p class="caption"><span class="caption-text">MindSpore Python API</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.html">mindspore</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.common.initializer.html">mindspore.common.initializer</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.communication.html">mindspore.communication</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.compression.html">mindspore.compression</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.context.html">mindspore.context</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.dataset.html">mindspore.dataset</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.dataset.audio.html">mindspore.dataset.audio</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.dataset.config.html">mindspore.dataset.config</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.dataset.text.html">mindspore.dataset.text</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.dataset.transforms.html">mindspore.dataset.transforms</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.dataset.vision.html">mindspore.dataset.vision</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.mindrecord.html">mindspore.mindrecord</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.nn.html">mindspore.nn</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.nn.probability.html">mindspore.nn.probability</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.nn.transformer.html">mindspore.nn.transformer</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.numpy.html">mindspore.numpy</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.ops.html">mindspore.ops</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.parallel.html">mindspore.parallel</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.parallel.nn.html">mindspore.parallel.nn</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.profiler.html">mindspore.profiler</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.scipy.html">mindspore.scipy</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.train.html">mindspore.train</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.boost.html">mindspore.boost</a></li>
</ul>
<p class="caption"><span class="caption-text">MindSpore C++ API</span></p>
<ul>
<li class="toctree-l1"><a class="reference external" href="https://www.mindspore.cn/lite/api/zh-CN/master/api_cpp/mindspore.html">MindSpore Lite↗</a></li>
</ul>

            
          
        </div>
        
      </div>
    </nav>

    <section data-toggle="wy-nav-shift" class="wy-nav-content-wrap">

      
      <nav class="wy-nav-top" aria-label="top navigation">
        
          <i data-toggle="wy-nav-top" class="fa fa-bars"></i>
          <a href="../../../../index.html">MindSpore</a>
        
      </nav>


      <div class="wy-nav-content">
        
        <div class="rst-content">
        
          

















<div role="navigation" aria-label="breadcrumbs navigation">

  <ul class="wy-breadcrumbs">
    
      <li><a href="../../../../index.html" class="icon icon-home"></a> &raquo;</li>
        
          <li><a href="../../../index.html">Module code</a> &raquo;</li>
        
      <li>mindspore.ops.composite.math_ops</li>
    
    
      <li class="wy-breadcrumbs-aside">
        
      </li>
    
  </ul>

  
  <hr/>
</div>
          <div role="main" class="document" itemscope="itemscope" itemtype="http://schema.org/Article">
           <div itemprop="articleBody">
            
  <h1>Source code for mindspore.ops.composite.math_ops</h1><div class="highlight"><pre>
<span></span><span class="c1"># Copyright 2020 Huawei Technologies Co., Ltd</span>
<span class="c1">#</span>
<span class="c1"># Licensed under the Apache License, Version 2.0 (the &quot;License&quot;);</span>
<span class="c1"># you may not use this file except in compliance with the License.</span>
<span class="c1"># You may obtain a copy of the License at</span>
<span class="c1">#</span>
<span class="c1"># http://www.apache.org/licenses/LICENSE-2.0</span>
<span class="c1">#</span>
<span class="c1"># Unless required by applicable law or agreed to in writing, software</span>
<span class="c1"># distributed under the License is distributed on an &quot;AS IS&quot; BASIS,</span>
<span class="c1"># WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.</span>
<span class="c1"># See the License for the specific language governing permissions and</span>
<span class="c1"># limitations under the License.</span>
<span class="c1"># ============================================================================</span>
<span class="sd">&quot;&quot;&quot;math Operations.&quot;&quot;&quot;</span>
<span class="kn">from</span> <span class="nn">itertools</span> <span class="kn">import</span> <span class="n">zip_longest</span>
<span class="kn">from</span> <span class="nn">collections</span> <span class="kn">import</span> <span class="n">deque</span>
<span class="kn">import</span> <span class="nn">numpy</span> <span class="k">as</span> <span class="nn">np</span>
<span class="kn">from</span> <span class="nn">mindspore.ops.composite.multitype_ops</span> <span class="kn">import</span> <span class="n">_constexpr_utils</span> <span class="k">as</span> <span class="n">const_utils</span>
<span class="kn">from</span> <span class="nn">mindspore.common</span> <span class="kn">import</span> <span class="n">dtype</span> <span class="k">as</span> <span class="n">mstype</span>
<span class="kn">from</span> <span class="nn">mindspore._checkparam</span> <span class="kn">import</span> <span class="n">Validator</span> <span class="k">as</span> <span class="n">validator</span>
<span class="kn">from</span> <span class="nn">mindspore.ops.operations</span> <span class="kn">import</span> <span class="n">_inner_ops</span> <span class="k">as</span> <span class="n">inner</span>
<span class="kn">from</span> <span class="nn">mindspore.ops.primitive</span> <span class="kn">import</span> <span class="n">constexpr</span>
<span class="kn">from</span> <span class="nn">mindspore.ops</span> <span class="kn">import</span> <span class="n">functional</span> <span class="k">as</span> <span class="n">F</span>
<span class="kn">from</span> <span class="nn">mindspore.ops.operations._inner_ops</span> <span class="kn">import</span> <span class="n">DynamicResizeNearestNeighbor</span>
<span class="kn">from</span> <span class="nn">..</span> <span class="kn">import</span> <span class="n">operations</span> <span class="k">as</span> <span class="n">P</span>


<span class="nd">@constexpr</span>
<span class="k">def</span> <span class="nf">_check_validate_axis</span><span class="p">(</span><span class="n">axis</span><span class="p">,</span> <span class="n">name</span><span class="p">):</span>
    <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">axis</span><span class="p">,</span> <span class="p">(</span><span class="nb">tuple</span><span class="p">,</span> <span class="nb">list</span><span class="p">)):</span>
        <span class="k">for</span> <span class="n">idx</span><span class="p">,</span> <span class="n">item</span> <span class="ow">in</span> <span class="nb">enumerate</span><span class="p">(</span><span class="n">axis</span><span class="p">):</span>
            <span class="n">validator</span><span class="o">.</span><span class="n">check_value_type</span><span class="p">(</span><span class="s2">&quot;axis[</span><span class="si">%d</span><span class="s2">]&quot;</span> <span class="o">%</span> <span class="n">idx</span><span class="p">,</span> <span class="n">item</span><span class="p">,</span> <span class="p">[</span><span class="nb">int</span><span class="p">],</span> <span class="n">name</span><span class="p">)</span>
    <span class="n">axis</span> <span class="o">=</span> <span class="n">validator</span><span class="o">.</span><span class="n">check_value_type</span><span class="p">(</span><span class="s1">&#39;axis&#39;</span><span class="p">,</span> <span class="n">axis</span><span class="p">,</span> <span class="p">[</span><span class="nb">int</span><span class="p">,</span> <span class="nb">tuple</span><span class="p">,</span> <span class="nb">list</span><span class="p">],</span> <span class="n">name</span><span class="p">)</span>
    <span class="k">return</span> <span class="n">axis</span>


<span class="nd">@constexpr</span>
<span class="k">def</span> <span class="nf">_check_validate_keepdims</span><span class="p">(</span><span class="n">keep_dims</span><span class="p">,</span> <span class="n">name</span><span class="p">):</span>
    <span class="n">keep_dims</span> <span class="o">=</span> <span class="n">validator</span><span class="o">.</span><span class="n">check_value_type</span><span class="p">(</span><span class="s1">&#39;keep_dims&#39;</span><span class="p">,</span> <span class="n">keep_dims</span><span class="p">,</span> <span class="p">[</span><span class="nb">bool</span><span class="p">],</span> <span class="n">name</span><span class="p">)</span>
    <span class="k">return</span> <span class="n">keep_dims</span>

<span class="nd">@constexpr</span>
<span class="k">def</span> <span class="nf">is_const</span><span class="p">(</span><span class="n">x</span><span class="p">):</span>
    <span class="k">return</span> <span class="n">x</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span>

<div class="viewcode-block" id="count_nonzero"><a class="viewcode-back" href="../../../../api_python/ops/mindspore.ops.count_nonzero.html#mindspore.ops.count_nonzero">[docs]</a><span class="k">def</span> <span class="nf">count_nonzero</span><span class="p">(</span><span class="n">x</span><span class="p">,</span> <span class="n">axis</span><span class="o">=</span><span class="p">(),</span> <span class="n">keep_dims</span><span class="o">=</span><span class="kc">False</span><span class="p">,</span> <span class="n">dtype</span><span class="o">=</span><span class="n">mstype</span><span class="o">.</span><span class="n">int32</span><span class="p">):</span>
    <span class="sa">r</span><span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Count number of nonzero elements across axis of input tensor</span>

<span class="sd">    Args:</span>
<span class="sd">        x (Tensor): Input data is used to count non-zero numbers.</span>
<span class="sd">          :math:`(N,*)` where :math:`*` means, any number of additional dimensions.</span>
<span class="sd">        axis (Union[int, tuple(int), list(int)]): The dimensions to reduce. Only constant value is allowed.</span>
<span class="sd">                                                  Default: (), reduce all dimensions.</span>
<span class="sd">        keep_dims (bool): If true, keep these reduced dimensions and the length is 1.</span>
<span class="sd">                          If false, don&#39;t keep these dimensions. Default: False.</span>
<span class="sd">        dtype (Union[Number, mindspore.bool\_]): The data type of the output tensor. Only constant value is allowed.</span>
<span class="sd">                                             Default: mindspore.int32</span>

<span class="sd">    Returns:</span>
<span class="sd">          Tensor, number of nonzero element. The data type is `dtype`.</span>

<span class="sd">    Supported Platforms:</span>
<span class="sd">        ``Ascend`` ``GPU`` ``CPU``</span>

<span class="sd">    Examples:</span>
<span class="sd">        &gt;&gt;&gt; from mindspore import Tensor, ops</span>
<span class="sd">        &gt;&gt;&gt; import numpy as np</span>
<span class="sd">        &gt;&gt;&gt; # case 1: each value specified.</span>
<span class="sd">        &gt;&gt;&gt; x = Tensor(np.array([[0, 1, 0], [1, 1, 0]]).astype(np.float32))</span>
<span class="sd">        &gt;&gt;&gt; nonzero_num = ops.count_nonzero(x=x, axis=[0, 1], keep_dims=True, dtype=mindspore.int32)</span>
<span class="sd">        &gt;&gt;&gt; print(nonzero_num)</span>
<span class="sd">        [[3]]</span>
<span class="sd">        &gt;&gt;&gt; # case 2: all value is default.</span>
<span class="sd">        &gt;&gt;&gt; nonzero_num = ops.count_nonzero(x=x)</span>
<span class="sd">        &gt;&gt;&gt; print(nonzero_num)</span>
<span class="sd">        3</span>
<span class="sd">        &gt;&gt;&gt; # case 3: axis value was specified 0.</span>
<span class="sd">        &gt;&gt;&gt; nonzero_num = ops.count_nonzero(x=x, axis=[0,])</span>
<span class="sd">        &gt;&gt;&gt; print(nonzero_num)</span>
<span class="sd">        [1 2 0]</span>
<span class="sd">        &gt;&gt;&gt; # case 4: axis value was specified 1.</span>
<span class="sd">        &gt;&gt;&gt; nonzero_num = ops.count_nonzero(x=x, axis=[1,])</span>
<span class="sd">        &gt;&gt;&gt; print(nonzero_num)</span>
<span class="sd">        [1 2]</span>
<span class="sd">        &gt;&gt;&gt; # case 5: keep_dims value was specified.</span>
<span class="sd">        &gt;&gt;&gt; nonzero_num = ops.count_nonzero(x=x,  keep_dims=True)</span>
<span class="sd">        &gt;&gt;&gt; print(nonzero_num)</span>
<span class="sd">        [[3]]</span>
<span class="sd">        &gt;&gt;&gt; # case 6: keep_dims and axis value was specified.</span>
<span class="sd">        &gt;&gt;&gt; nonzero_num = ops.count_nonzero(x=x, axis=[0,], keep_dims=True)</span>
<span class="sd">        &gt;&gt;&gt; print(nonzero_num)</span>
<span class="sd">        [[1 2 0]]</span>
<span class="sd">    &quot;&quot;&quot;</span>

    <span class="n">const_utils</span><span class="o">.</span><span class="n">check_type_valid</span><span class="p">(</span><span class="n">F</span><span class="o">.</span><span class="n">dtype</span><span class="p">(</span><span class="n">x</span><span class="p">),</span> <span class="n">mstype</span><span class="o">.</span><span class="n">number_type</span><span class="p">,</span> <span class="s1">&#39;input x&#39;</span><span class="p">)</span>
    <span class="n">axis</span> <span class="o">=</span> <span class="n">_check_validate_axis</span><span class="p">(</span><span class="n">axis</span><span class="p">,</span> <span class="s2">&quot;count_nonzero&quot;</span><span class="p">)</span>
    <span class="n">keep_dims</span> <span class="o">=</span> <span class="n">_check_validate_keepdims</span><span class="p">(</span><span class="n">keep_dims</span><span class="p">,</span> <span class="s2">&quot;count_nonzero&quot;</span><span class="p">)</span>
    <span class="n">const_utils</span><span class="o">.</span><span class="n">check_type_valid</span><span class="p">(</span><span class="n">dtype</span><span class="p">,</span> <span class="n">mstype</span><span class="o">.</span><span class="n">number_type</span> <span class="o">+</span> <span class="p">(</span><span class="n">mstype</span><span class="o">.</span><span class="n">bool_</span><span class="p">,),</span> <span class="s1">&#39;dtype&#39;</span><span class="p">)</span>

    <span class="n">not_equal</span> <span class="o">=</span> <span class="n">P</span><span class="o">.</span><span class="n">NotEqual</span><span class="p">()</span>
    <span class="n">cast</span> <span class="o">=</span> <span class="n">P</span><span class="o">.</span><span class="n">Cast</span><span class="p">()</span>
    <span class="n">reduce_sum</span> <span class="o">=</span> <span class="n">P</span><span class="o">.</span><span class="n">ReduceSum</span><span class="p">(</span><span class="n">keep_dims</span><span class="p">)</span>
    <span class="n">nonzero_bool</span> <span class="o">=</span> <span class="n">not_equal</span><span class="p">(</span><span class="n">x</span><span class="p">,</span> <span class="mi">0</span><span class="p">)</span>
    <span class="c1"># ReduceSum only support float16 or float32 tensor.</span>
    <span class="n">nonzero_val</span> <span class="o">=</span> <span class="n">cast</span><span class="p">(</span><span class="n">nonzero_bool</span><span class="p">,</span> <span class="n">mstype</span><span class="o">.</span><span class="n">float32</span><span class="p">)</span>
    <span class="n">nonzero_num</span> <span class="o">=</span> <span class="n">cast</span><span class="p">(</span><span class="n">reduce_sum</span><span class="p">(</span><span class="n">nonzero_val</span><span class="p">,</span> <span class="n">axis</span><span class="p">),</span> <span class="n">dtype</span><span class="p">)</span>

    <span class="k">return</span> <span class="n">nonzero_num</span></div>


<span class="nd">@constexpr</span>
<span class="k">def</span> <span class="nf">_int_to_tuple_conv</span><span class="p">(</span><span class="n">axes</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Converts ints to tuples in input axes, expected by most validation checks.</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="k">for</span> <span class="n">x</span> <span class="ow">in</span> <span class="p">[</span><span class="mi">0</span><span class="p">,</span> <span class="mi">1</span><span class="p">]:</span>
        <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">axes</span><span class="p">[</span><span class="n">x</span><span class="p">],</span> <span class="nb">int</span><span class="p">):</span>
            <span class="n">axes</span><span class="p">[</span><span class="n">x</span><span class="p">]</span> <span class="o">=</span> <span class="p">(</span><span class="n">axes</span><span class="p">[</span><span class="n">x</span><span class="p">],)</span>
    <span class="k">return</span> <span class="n">axes</span>


<span class="nd">@constexpr</span>
<span class="k">def</span> <span class="nf">_check_axes</span><span class="p">(</span><span class="n">axes</span><span class="p">,</span> <span class="n">prim_name</span><span class="o">=</span><span class="kc">None</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Check for validity and type of axes passed to function.</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="n">msg_prefix</span> <span class="o">=</span> <span class="sa">f</span><span class="s2">&quot;For &#39;</span><span class="si">{</span><span class="n">prim_name</span><span class="si">}</span><span class="s2">&#39;, the&quot;</span> <span class="k">if</span> <span class="n">prim_name</span> <span class="k">else</span> <span class="s2">&quot;The&quot;</span>
    <span class="n">validator</span><span class="o">.</span><span class="n">check_value_type</span><span class="p">(</span><span class="s1">&#39;axes&#39;</span><span class="p">,</span> <span class="n">axes</span><span class="p">,</span> <span class="p">[</span><span class="nb">int</span><span class="p">,</span> <span class="nb">tuple</span><span class="p">,</span> <span class="nb">list</span><span class="p">],</span> <span class="s2">&quot;tensor dot&quot;</span><span class="p">)</span>
    <span class="k">if</span> <span class="ow">not</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">axes</span><span class="p">,</span> <span class="nb">int</span><span class="p">):</span>
        <span class="n">axes</span> <span class="o">=</span> <span class="nb">list</span><span class="p">(</span><span class="n">axes</span><span class="p">)</span>  <span class="c1"># to avoid immutability issues</span>
        <span class="k">if</span> <span class="nb">len</span><span class="p">(</span><span class="n">axes</span><span class="p">)</span> <span class="o">!=</span> <span class="mi">2</span><span class="p">:</span>
            <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="sa">f</span><span class="s2">&quot;</span><span class="si">{</span><span class="n">msg_prefix</span><span class="si">}</span><span class="s2"> dimension of &#39;axes&#39; should be 2, but got &#39;axes&#39;: </span><span class="si">{</span><span class="n">axes</span><span class="si">}</span><span class="s2">.&quot;</span><span class="p">)</span>
        <span class="n">axes</span> <span class="o">=</span> <span class="n">_int_to_tuple_conv</span><span class="p">(</span><span class="n">axes</span><span class="p">)</span>  <span class="c1"># convert before length checks</span>
        <span class="k">if</span> <span class="nb">len</span><span class="p">(</span><span class="n">axes</span><span class="p">[</span><span class="mi">0</span><span class="p">])</span> <span class="o">!=</span> <span class="nb">len</span><span class="p">(</span><span class="n">axes</span><span class="p">[</span><span class="mi">1</span><span class="p">]):</span>
            <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="sa">f</span><span class="s2">&quot;</span><span class="si">{</span><span class="n">msg_prefix</span><span class="si">}</span><span class="s2"> first and second dim of &#39;axes&#39; have to be the same size/length, &quot;</span>
                             <span class="sa">f</span><span class="s2">&quot;but got &#39;axes&#39;: </span><span class="si">{</span><span class="n">axes</span><span class="si">}</span><span class="s2">.&quot;</span><span class="p">)</span>
        <span class="k">if</span> <span class="nb">len</span><span class="p">(</span><span class="n">axes</span><span class="p">[</span><span class="mi">0</span><span class="p">])</span> <span class="o">!=</span> <span class="nb">len</span><span class="p">(</span><span class="nb">set</span><span class="p">(</span><span class="n">axes</span><span class="p">[</span><span class="mi">0</span><span class="p">]))</span> <span class="ow">or</span> <span class="nb">len</span><span class="p">(</span><span class="n">axes</span><span class="p">[</span><span class="mi">1</span><span class="p">])</span> <span class="o">!=</span> <span class="nb">len</span><span class="p">(</span><span class="nb">set</span><span class="p">(</span><span class="n">axes</span><span class="p">[</span><span class="mi">1</span><span class="p">])):</span>
            <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="sa">f</span><span class="s2">&quot;</span><span class="si">{</span><span class="n">msg_prefix</span><span class="si">}</span><span class="s2"> &#39;axes&#39; cannot have duplicating values, but got </span><span class="si">{</span><span class="n">axes</span><span class="si">}</span><span class="s2">.&quot;</span><span class="p">)</span>
    <span class="k">return</span> <span class="n">axes</span>


<span class="nd">@constexpr</span>
<span class="k">def</span> <span class="nf">_typecheck_input</span><span class="p">(</span><span class="n">x1_type</span><span class="p">,</span> <span class="n">x2_type</span><span class="p">,</span> <span class="n">prim_name</span><span class="o">=</span><span class="kc">None</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Check input tensor types to be valid and confirm they are the same type.</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="n">msg_prefix</span> <span class="o">=</span> <span class="sa">f</span><span class="s2">&quot;For &#39;</span><span class="si">{</span><span class="n">prim_name</span><span class="si">}</span><span class="s2">&#39;, the&quot;</span> <span class="k">if</span> <span class="n">prim_name</span> <span class="k">else</span> <span class="s2">&quot;The&quot;</span>
    <span class="n">const_utils</span><span class="o">.</span><span class="n">check_type_valid</span><span class="p">(</span><span class="n">x1_type</span><span class="p">,</span> <span class="p">[</span><span class="n">mstype</span><span class="o">.</span><span class="n">float32</span><span class="p">,</span> <span class="n">mstype</span><span class="o">.</span><span class="n">float16</span><span class="p">],</span> <span class="s1">&#39;x1&#39;</span><span class="p">)</span>
    <span class="n">const_utils</span><span class="o">.</span><span class="n">check_type_valid</span><span class="p">(</span><span class="n">x2_type</span><span class="p">,</span> <span class="p">[</span><span class="n">mstype</span><span class="o">.</span><span class="n">float32</span><span class="p">,</span> <span class="n">mstype</span><span class="o">.</span><span class="n">float16</span><span class="p">],</span> <span class="s1">&#39;x2&#39;</span><span class="p">)</span>
    <span class="k">if</span> <span class="n">x1_type</span> <span class="o">!=</span> <span class="n">x2_type</span><span class="p">:</span>
        <span class="k">raise</span> <span class="ne">TypeError</span><span class="p">(</span><span class="sa">f</span><span class="s2">&quot;</span><span class="si">{</span><span class="n">msg_prefix</span><span class="si">}</span><span class="s2"> inputs must be the same type, but got x1_type: </span><span class="si">{</span><span class="n">x1_type</span><span class="si">}</span><span class="s2"> &quot;</span>
                        <span class="sa">f</span><span class="s2">&quot;and x2_type: </span><span class="si">{</span><span class="n">x2_type</span><span class="si">}</span><span class="s2">.&quot;</span><span class="p">)</span>


<span class="nd">@constexpr</span>
<span class="k">def</span> <span class="nf">_axes_int_check</span><span class="p">(</span><span class="n">x1_shape</span><span class="p">,</span> <span class="n">x2_shape</span><span class="p">,</span> <span class="n">axes</span><span class="p">,</span> <span class="n">prim_name</span><span class="o">=</span><span class="kc">None</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Convert from single int axes to 2d tuple if required</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="n">msg_prefix</span> <span class="o">=</span> <span class="sa">f</span><span class="s2">&quot;For &#39;</span><span class="si">{</span><span class="n">prim_name</span><span class="si">}</span><span class="s2">&#39;, the&quot;</span> <span class="k">if</span> <span class="n">prim_name</span> <span class="k">else</span> <span class="s2">&quot;The&quot;</span>
    <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">axes</span><span class="p">,</span> <span class="nb">int</span><span class="p">):</span>
        <span class="k">if</span> <span class="n">axes</span> <span class="o">&lt;</span> <span class="mi">0</span><span class="p">:</span>
            <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="sa">f</span><span class="s2">&quot;</span><span class="si">{</span><span class="n">msg_prefix</span><span class="si">}</span><span class="s2"> &#39;axes&#39; must be at least 0, but got </span><span class="si">{</span><span class="n">axes</span><span class="si">}</span><span class="s2">.&quot;</span><span class="p">)</span>
        <span class="k">if</span> <span class="n">axes</span> <span class="o">==</span> <span class="mi">0</span><span class="p">:</span>
            <span class="c1"># outer product, no input validation required</span>
            <span class="k">return</span> <span class="p">[],</span> <span class="p">[]</span>
        <span class="k">if</span> <span class="n">axes</span> <span class="o">&gt;</span> <span class="nb">len</span><span class="p">(</span><span class="n">x1_shape</span><span class="p">)</span> <span class="ow">or</span> <span class="n">axes</span> <span class="o">&gt;</span> <span class="nb">len</span><span class="p">(</span><span class="n">x2_shape</span><span class="p">):</span>
            <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="sa">f</span><span class="s2">&quot;</span><span class="si">{</span><span class="n">msg_prefix</span><span class="si">}</span><span class="s2"> &#39;axes&#39; cannot be greater than the length of &#39;x1_shape&#39; and &#39;x2_shape&#39;, &quot;</span>
                             <span class="sa">f</span><span class="s2">&quot;but got &#39;axes&#39;: </span><span class="si">{</span><span class="n">axes</span><span class="si">}</span><span class="s2">, &#39;x1_shape&#39;: </span><span class="si">{</span><span class="n">x1_shape</span><span class="si">}</span><span class="s2">, &#39;x2_shape&#39;: </span><span class="si">{</span><span class="n">x2_shape</span><span class="si">}</span><span class="s2">.&quot;</span><span class="p">)</span>
        <span class="n">x1_ind</span> <span class="o">=</span> <span class="nb">tuple</span><span class="p">(</span><span class="nb">range</span><span class="p">(</span><span class="nb">len</span><span class="p">(</span><span class="n">x1_shape</span><span class="p">))[</span><span class="o">-</span><span class="mi">1</span> <span class="o">*</span> <span class="n">axes</span><span class="p">:])</span>
        <span class="n">x2_ind</span> <span class="o">=</span> <span class="nb">tuple</span><span class="p">(</span><span class="nb">range</span><span class="p">(</span><span class="nb">len</span><span class="p">(</span><span class="n">x2_shape</span><span class="p">))[:</span><span class="n">axes</span><span class="p">])</span>
        <span class="n">axes</span> <span class="o">=</span> <span class="nb">tuple</span><span class="p">((</span><span class="n">x1_ind</span><span class="p">,</span> <span class="n">x2_ind</span><span class="p">))</span>
        <span class="n">axes</span> <span class="o">=</span> <span class="n">_int_to_tuple_conv</span><span class="p">(</span><span class="n">axes</span><span class="p">)</span>
    <span class="k">return</span> <span class="n">axes</span>


<span class="nd">@constexpr</span>
<span class="k">def</span> <span class="nf">_validate_axes</span><span class="p">(</span><span class="n">x1_shape</span><span class="p">,</span> <span class="n">x2_shape</span><span class="p">,</span> <span class="n">axes</span><span class="p">,</span> <span class="n">prim_name</span><span class="o">=</span><span class="kc">None</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Checks for axes having the correct length according to input, for any value in axis</span>
<span class="sd">    being out of range with given shape and also checking for compatible axes values</span>
<span class="sd">    with given inputs.</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="n">msg_prefix</span> <span class="o">=</span> <span class="sa">f</span><span class="s2">&quot;For &#39;</span><span class="si">{</span><span class="n">prim_name</span><span class="si">}</span><span class="s2">&#39;, the&quot;</span> <span class="k">if</span> <span class="n">prim_name</span> <span class="k">else</span> <span class="s2">&quot;The&quot;</span>
    <span class="n">shapes</span> <span class="o">=</span> <span class="p">[</span><span class="n">x1_shape</span><span class="p">,</span> <span class="n">x2_shape</span><span class="p">]</span>

    <span class="c1"># axis length check</span>
    <span class="k">for</span> <span class="n">ix_input</span><span class="p">,</span> <span class="n">x_axes</span> <span class="ow">in</span> <span class="nb">enumerate</span><span class="p">(</span><span class="n">axes</span><span class="p">):</span>
        <span class="n">axes_len</span> <span class="o">=</span> <span class="nb">len</span><span class="p">(</span><span class="n">x_axes</span><span class="p">)</span>
        <span class="n">shape_dim_len</span> <span class="o">=</span> <span class="nb">len</span><span class="p">(</span><span class="n">shapes</span><span class="p">[</span><span class="n">ix_input</span><span class="p">])</span>
        <span class="k">if</span> <span class="n">axes_len</span> <span class="o">&gt;</span> <span class="n">shape_dim_len</span><span class="p">:</span>
            <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="sa">f</span><span class="s2">&quot;</span><span class="si">{</span><span class="n">msg_prefix</span><span class="si">}</span><span class="s2"> length of element </span><span class="si">{</span><span class="n">x_axes</span><span class="si">}</span><span class="s2"> in &#39;axes&#39; should be less than or equal to &quot;</span>
                             <span class="sa">f</span><span class="s2">&quot;</span><span class="si">{</span><span class="n">shape_dim_len</span><span class="si">}</span><span class="s2">, but got </span><span class="si">{</span><span class="n">axes_len</span><span class="si">}</span><span class="s2">.&quot;</span><span class="p">)</span>

    <span class="c1"># axis values range check</span>
    <span class="k">for</span> <span class="n">ix_input</span><span class="p">,</span> <span class="n">x_axes</span> <span class="ow">in</span> <span class="nb">enumerate</span><span class="p">(</span><span class="n">axes</span><span class="p">):</span>
        <span class="n">comp_shape</span> <span class="o">=</span> <span class="n">shapes</span><span class="p">[</span><span class="n">ix_input</span><span class="p">]</span>
        <span class="n">max_val</span> <span class="o">=</span> <span class="nb">len</span><span class="p">(</span><span class="n">comp_shape</span><span class="p">)</span> <span class="o">-</span> <span class="mi">1</span>
        <span class="n">min_val</span> <span class="o">=</span> <span class="o">-</span><span class="mi">1</span> <span class="o">*</span> <span class="nb">len</span><span class="p">(</span><span class="n">comp_shape</span><span class="p">)</span>
        <span class="k">for</span> <span class="n">_</span><span class="p">,</span> <span class="n">x_value</span> <span class="ow">in</span> <span class="nb">enumerate</span><span class="p">(</span><span class="n">x_axes</span><span class="p">):</span>
            <span class="k">if</span> <span class="ow">not</span> <span class="n">min_val</span> <span class="o">&lt;=</span> <span class="n">x_value</span> <span class="o">&lt;=</span> <span class="n">max_val</span><span class="p">:</span>
                <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="sa">f</span><span class="s2">&quot;</span><span class="si">{</span><span class="n">msg_prefix</span><span class="si">}</span><span class="s2"> value in &#39;axes&#39; should be in range: [</span><span class="si">{</span><span class="n">min_val</span><span class="si">}</span><span class="s2">, </span><span class="si">{</span><span class="n">max_val</span><span class="si">}</span><span class="s2">], &quot;</span>
                                 <span class="sa">f</span><span class="s2">&quot;but got </span><span class="si">{</span><span class="n">x_value</span><span class="si">}</span><span class="s2">.&quot;</span><span class="p">)</span>

    <span class="c1"># check axis value with input shape - both ways for axis valid</span>
    <span class="n">invalid_a</span> <span class="o">=</span> <span class="kc">False</span>
    <span class="n">invalid_b</span> <span class="o">=</span> <span class="kc">False</span>
    <span class="k">for</span> <span class="n">i</span> <span class="ow">in</span> <span class="nb">range</span><span class="p">(</span><span class="nb">len</span><span class="p">(</span><span class="n">axes</span><span class="p">[</span><span class="mi">0</span><span class="p">])):</span>  <span class="c1"># sizes already validated</span>
        <span class="k">if</span> <span class="n">x1_shape</span><span class="p">[</span><span class="n">axes</span><span class="p">[</span><span class="mi">0</span><span class="p">][</span><span class="n">i</span><span class="p">]]</span> <span class="o">!=</span> <span class="n">x2_shape</span><span class="p">[</span><span class="n">axes</span><span class="p">[</span><span class="mi">1</span><span class="p">][</span><span class="n">i</span><span class="p">]]:</span>
            <span class="n">invalid_a</span> <span class="o">=</span> <span class="kc">True</span>
        <span class="k">if</span> <span class="n">x1_shape</span><span class="p">[</span><span class="n">axes</span><span class="p">[</span><span class="mi">0</span><span class="p">][</span><span class="n">i</span><span class="p">]]</span> <span class="o">!=</span> <span class="n">x2_shape</span><span class="p">[</span><span class="n">axes</span><span class="p">[</span><span class="mi">1</span><span class="p">][</span><span class="nb">len</span><span class="p">(</span><span class="n">axes</span><span class="p">[</span><span class="mi">0</span><span class="p">])</span> <span class="o">-</span> <span class="mi">1</span> <span class="o">-</span> <span class="n">i</span><span class="p">]]:</span>
            <span class="n">invalid_b</span> <span class="o">=</span> <span class="kc">True</span>
    <span class="k">if</span> <span class="n">invalid_a</span> <span class="ow">and</span> <span class="n">invalid_b</span><span class="p">:</span>
        <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="sa">f</span><span class="s2">&quot;</span><span class="si">{</span><span class="n">msg_prefix</span><span class="si">}</span><span class="s2"> &#39;i&#39; should exist such that &#39;x1_shape[axes[0][i]]&#39; is equal to &quot;</span>
                         <span class="sa">f</span><span class="s2">&quot;&#39;x2_shape[axes[1][i]]&#39; or &#39;x2_shape[axes[1][len(axes[0])-1-i]]&#39;, but got &quot;</span>
                         <span class="sa">f</span><span class="s2">&quot;&#39;x1_shape&#39;: </span><span class="si">{</span><span class="n">x1_shape</span><span class="si">}</span><span class="s2">, &#39;x2_shape&#39;: </span><span class="si">{</span><span class="n">x2_shape</span><span class="si">}</span><span class="s2">, &#39;axes&#39;: </span><span class="si">{</span><span class="n">axes</span><span class="si">}</span><span class="s2">.&quot;</span><span class="p">)</span>


<span class="nd">@constexpr</span>
<span class="k">def</span> <span class="nf">_calc_new_shape</span><span class="p">(</span><span class="n">shape</span><span class="p">,</span> <span class="n">axes</span><span class="p">,</span> <span class="n">position</span><span class="o">=</span><span class="mi">0</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Calculate transpose and reshape parameters for input transformations,</span>
<span class="sd">    &#39;position&#39; refers to whether tensor is first or second in the op.</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="n">contraction_axes</span> <span class="o">=</span> <span class="nb">tuple</span><span class="p">(</span><span class="n">i</span> <span class="k">if</span> <span class="n">i</span> <span class="o">&gt;=</span> <span class="mi">0</span> <span class="k">else</span> <span class="n">i</span> <span class="o">+</span> <span class="nb">len</span><span class="p">(</span><span class="n">shape</span><span class="p">)</span> <span class="k">for</span> <span class="n">i</span> <span class="ow">in</span> <span class="n">axes</span><span class="p">[</span><span class="n">position</span><span class="p">])</span>
    <span class="n">prod_contraction</span> <span class="o">=</span> <span class="nb">int</span><span class="p">(</span><span class="n">np</span><span class="o">.</span><span class="n">prod</span><span class="p">([</span><span class="n">shape</span><span class="p">[</span><span class="n">i</span><span class="p">]</span> <span class="k">for</span> <span class="n">i</span> <span class="ow">in</span> <span class="n">contraction_axes</span><span class="p">]))</span>
    <span class="n">free_axes</span> <span class="o">=</span> <span class="nb">tuple</span><span class="p">(</span><span class="n">i</span> <span class="k">for</span> <span class="n">i</span> <span class="ow">in</span> <span class="nb">range</span><span class="p">(</span><span class="nb">len</span><span class="p">(</span><span class="n">shape</span><span class="p">))</span> <span class="k">if</span> <span class="n">i</span> <span class="ow">not</span> <span class="ow">in</span> <span class="n">contraction_axes</span><span class="p">)</span>
    <span class="n">free_dims</span> <span class="o">=</span> <span class="nb">tuple</span><span class="p">(</span><span class="n">shape</span><span class="p">[</span><span class="n">i</span><span class="p">]</span> <span class="k">for</span> <span class="n">i</span> <span class="ow">in</span> <span class="n">free_axes</span><span class="p">)</span>
    <span class="n">prod_free</span> <span class="o">=</span> <span class="nb">int</span><span class="p">(</span><span class="n">np</span><span class="o">.</span><span class="n">prod</span><span class="p">(</span><span class="n">free_dims</span><span class="p">))</span>

    <span class="n">transpose_perm</span> <span class="o">=</span> <span class="n">contraction_axes</span> <span class="o">+</span> <span class="n">free_axes</span> <span class="k">if</span> <span class="n">position</span> <span class="k">else</span> <span class="n">free_axes</span> <span class="o">+</span> <span class="n">contraction_axes</span>
    <span class="n">new_shape</span> <span class="o">=</span> <span class="p">(</span><span class="n">prod_contraction</span><span class="p">,</span> <span class="n">prod_free</span><span class="p">)</span> <span class="k">if</span> <span class="n">position</span> <span class="k">else</span> <span class="p">(</span><span class="n">prod_free</span><span class="p">,</span> <span class="n">prod_contraction</span><span class="p">)</span>
    <span class="k">return</span> <span class="n">new_shape</span><span class="p">,</span> <span class="n">transpose_perm</span><span class="p">,</span> <span class="n">free_dims</span>


<div class="viewcode-block" id="tensor_dot"><a class="viewcode-back" href="../../../../api_python/ops/mindspore.ops.tensor_dot.html#mindspore.ops.tensor_dot">[docs]</a><span class="k">def</span> <span class="nf">tensor_dot</span><span class="p">(</span><span class="n">x1</span><span class="p">,</span> <span class="n">x2</span><span class="p">,</span> <span class="n">axes</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Computation of Tensor contraction on arbitrary axes between tensors `a` and `b`.</span>

<span class="sd">    Contraction allows for the summation of products of elements of `a` and `b` on specified axes.</span>
<span class="sd">    The same number of axes must be specified for both x1 and x2, and values must be within range</span>
<span class="sd">    of number of dims of both `a` and `b`.</span>

<span class="sd">    Selected dims in both inputs must also match.</span>

<span class="sd">    axes = 0 leads to outer product</span>
<span class="sd">    axes = 1 leads to normal matrix multiplication when inputs both 2D.</span>
<span class="sd">    axes = 1 is the same as axes = ((1,),(0,)) where both `a` and `b` are 2D.</span>
<span class="sd">    axes = 2 is the same as axes = ((1,2),(0,1)) where both `a` and `b` are 3D.</span>

<span class="sd">    Args:</span>
<span class="sd">        x1 (Tensor): First tensor in tensor_dot with datatype float16 or float32</span>
<span class="sd">        x2 (Tensor): Second tensor in tensor_dot with datatype float16 or float32</span>
<span class="sd">        axes (Union[int, tuple(int), tuple(tuple(int)), list(list(int))]): Single value or</span>
<span class="sd">            tuple/list of length 2 with dimensions specified for `a` and `b` each. If single value `N` passed,</span>
<span class="sd">            automatically picks up last N dims from `a` input shape and first N dims from `b` input shape in order</span>
<span class="sd">            as axes for each respectively.</span>

<span class="sd">    Inputs:</span>
<span class="sd">        - **x1** (Tensor) - First tensor in tensor_dot with datatype float16 or float32</span>
<span class="sd">        - **x2** (Tensor) - Second tensor in tensor_dot with datatype float16 or float32</span>
<span class="sd">        - **axes** (Union[int, tuple(int), tuple(tuple(int)), list(list(int))]) - Single value or</span>
<span class="sd">          tuple/list of length 2 with dimensions specified for `a` and `b` each. If single value `N` passed,</span>
<span class="sd">          automatically picks up last N dims from `a` input shape and first N dims from `b` input shape in order</span>
<span class="sd">          as axes for each respectively.</span>

<span class="sd">    Outputs:</span>
<span class="sd">        Tensor, the shape of the output tensor is :math:`(N + M)`. Where :math:`N` and :math:`M` are the free axes not</span>
<span class="sd">        contracted in both inputs</span>

<span class="sd">    Raises:</span>
<span class="sd">        TypeError: If `x1` or `x2` is not a Tensor.</span>
<span class="sd">        TypeError: If `axes` is not one of the following: int, tuple, list.</span>

<span class="sd">    Supported Platforms:</span>
<span class="sd">        ``Ascend`` ``GPU`` ``CPU``</span>

<span class="sd">    Examples:</span>
<span class="sd">        &gt;&gt;&gt; from mindspore import Tensor, ops</span>
<span class="sd">        &gt;&gt;&gt; import mindspore</span>
<span class="sd">        &gt;&gt;&gt; import numpy as np</span>
<span class="sd">        &gt;&gt;&gt; input_x1 = Tensor(np.ones(shape=[1, 2, 3]), mindspore.float32)</span>
<span class="sd">        &gt;&gt;&gt; input_x2 = Tensor(np.ones(shape=[3, 1, 2]), mindspore.float32)</span>
<span class="sd">        &gt;&gt;&gt; output = ops.tensor_dot(input_x1, input_x2, ((0,1),(1,2)))</span>
<span class="sd">        &gt;&gt;&gt; print(output)</span>
<span class="sd">        [[2. 2. 2]</span>
<span class="sd">         [2. 2. 2]</span>
<span class="sd">         [2. 2. 2]]</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="n">shape_op</span> <span class="o">=</span> <span class="n">P</span><span class="o">.</span><span class="n">Shape</span><span class="p">()</span>
    <span class="n">reshape_op</span> <span class="o">=</span> <span class="n">P</span><span class="o">.</span><span class="n">Reshape</span><span class="p">()</span>
    <span class="n">transpose_op</span> <span class="o">=</span> <span class="n">P</span><span class="o">.</span><span class="n">Transpose</span><span class="p">()</span>
    <span class="n">matmul_op</span> <span class="o">=</span> <span class="n">P</span><span class="o">.</span><span class="n">MatMul</span><span class="p">(</span><span class="kc">False</span><span class="p">,</span> <span class="kc">False</span><span class="p">)</span>
    <span class="c1"># input validity checks</span>
    <span class="n">x1_shape</span> <span class="o">=</span> <span class="n">shape_op</span><span class="p">(</span><span class="n">x1</span><span class="p">)</span>
    <span class="n">x2_shape</span> <span class="o">=</span> <span class="n">shape_op</span><span class="p">(</span><span class="n">x2</span><span class="p">)</span>
    <span class="n">x1_type</span> <span class="o">=</span> <span class="n">F</span><span class="o">.</span><span class="n">dtype</span><span class="p">(</span><span class="n">x1</span><span class="p">)</span>
    <span class="n">x2_type</span> <span class="o">=</span> <span class="n">F</span><span class="o">.</span><span class="n">dtype</span><span class="p">(</span><span class="n">x2</span><span class="p">)</span>
    <span class="n">axes</span> <span class="o">=</span> <span class="n">_check_axes</span><span class="p">(</span><span class="n">axes</span><span class="p">,</span> <span class="s1">&#39;tensor_dot&#39;</span><span class="p">)</span>
    <span class="n">_typecheck_input</span><span class="p">(</span><span class="n">x1_type</span><span class="p">,</span> <span class="n">x2_type</span><span class="p">,</span> <span class="s1">&#39;tensor_dot&#39;</span><span class="p">)</span>
    <span class="c1"># input compatibility check &amp; axes format update</span>
    <span class="n">axes</span> <span class="o">=</span> <span class="n">_axes_int_check</span><span class="p">(</span><span class="n">x1_shape</span><span class="p">,</span> <span class="n">x2_shape</span><span class="p">,</span> <span class="n">axes</span><span class="p">,</span> <span class="s1">&#39;tensor_dot&#39;</span><span class="p">)</span>
    <span class="n">_validate_axes</span><span class="p">(</span><span class="n">x1_shape</span><span class="p">,</span> <span class="n">x2_shape</span><span class="p">,</span> <span class="n">axes</span><span class="p">,</span> <span class="s1">&#39;tensor_dot&#39;</span><span class="p">)</span>
    <span class="n">x1_reshape_fwd</span><span class="p">,</span> <span class="n">x1_transpose_fwd</span><span class="p">,</span> <span class="n">x1_ret</span> <span class="o">=</span> <span class="n">_calc_new_shape</span><span class="p">(</span><span class="n">x1_shape</span><span class="p">,</span> <span class="n">axes</span><span class="p">,</span> <span class="mi">0</span><span class="p">)</span>
    <span class="n">x2_reshape_fwd</span><span class="p">,</span> <span class="n">x2_transpose_fwd</span><span class="p">,</span> <span class="n">x2_ret</span> <span class="o">=</span> <span class="n">_calc_new_shape</span><span class="p">(</span><span class="n">x2_shape</span><span class="p">,</span> <span class="n">axes</span><span class="p">,</span> <span class="mi">1</span><span class="p">)</span>
    <span class="n">output_shape</span> <span class="o">=</span> <span class="n">x1_ret</span> <span class="o">+</span> <span class="n">x2_ret</span>  <span class="c1"># combine free axes from both inputs</span>
    <span class="c1"># run tensor_dot op</span>
    <span class="n">x1_transposed</span> <span class="o">=</span> <span class="n">transpose_op</span><span class="p">(</span><span class="n">x1</span><span class="p">,</span> <span class="n">x1_transpose_fwd</span><span class="p">)</span>
    <span class="n">x2_transposed</span> <span class="o">=</span> <span class="n">transpose_op</span><span class="p">(</span><span class="n">x2</span><span class="p">,</span> <span class="n">x2_transpose_fwd</span><span class="p">)</span>
    <span class="n">x1_reshaped</span> <span class="o">=</span> <span class="n">reshape_op</span><span class="p">(</span><span class="n">x1_transposed</span><span class="p">,</span> <span class="n">x1_reshape_fwd</span><span class="p">)</span>
    <span class="n">x2_reshaped</span> <span class="o">=</span> <span class="n">reshape_op</span><span class="p">(</span><span class="n">x2_transposed</span><span class="p">,</span> <span class="n">x2_reshape_fwd</span><span class="p">)</span>
    <span class="n">mul_result</span> <span class="o">=</span> <span class="n">matmul_op</span><span class="p">(</span><span class="n">x1_reshaped</span><span class="p">,</span> <span class="n">x2_reshaped</span><span class="p">)</span>
    <span class="n">final_result</span> <span class="o">=</span> <span class="n">reshape_op</span><span class="p">(</span><span class="n">mul_result</span><span class="p">,</span> <span class="n">output_shape</span><span class="p">)</span>
    <span class="k">return</span> <span class="n">final_result</span></div>


<span class="nd">@constexpr</span>
<span class="k">def</span> <span class="nf">_check_invalid_input</span><span class="p">(</span><span class="n">x1_shape</span><span class="p">,</span> <span class="n">x2_shape</span><span class="p">,</span> <span class="n">prim_name</span><span class="o">=</span><span class="kc">None</span><span class="p">):</span>
    <span class="n">msg_prefix</span> <span class="o">=</span> <span class="sa">f</span><span class="s2">&quot;For &#39;</span><span class="si">{</span><span class="n">prim_name</span><span class="si">}</span><span class="s2">&#39;, the&quot;</span> <span class="k">if</span> <span class="n">prim_name</span> <span class="k">else</span> <span class="s2">&quot;The&quot;</span>
    <span class="k">if</span> <span class="nb">len</span><span class="p">(</span><span class="n">x1_shape</span><span class="p">)</span> <span class="o">&lt;</span> <span class="mi">2</span> <span class="ow">or</span> <span class="nb">len</span><span class="p">(</span><span class="n">x2_shape</span><span class="p">)</span> <span class="o">&lt;</span> <span class="mi">2</span><span class="p">:</span>
        <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="sa">f</span><span class="s2">&quot;</span><span class="si">{</span><span class="n">msg_prefix</span><span class="si">}</span><span class="s2"> inputs x1, x2 should have &#39;dimension &gt;= 2&#39;,&quot;</span>
                         <span class="sa">f</span><span class="s2">&quot;but got &#39;len(x1_shape)&#39;: (</span><span class="si">{</span><span class="nb">len</span><span class="p">(</span><span class="n">x1_shape</span><span class="p">)</span><span class="si">}</span><span class="s2">) and &#39;len(x2_shape)&#39;: (</span><span class="si">{</span><span class="nb">len</span><span class="p">(</span><span class="n">x2_shape</span><span class="p">)</span><span class="si">}</span><span class="s2">).&quot;</span><span class="p">)</span>


<span class="nd">@constexpr</span>
<span class="k">def</span> <span class="nf">_typecheck_input_dot</span><span class="p">(</span><span class="n">x1_type</span><span class="p">,</span> <span class="n">x2_type</span><span class="p">,</span> <span class="n">prim_name</span><span class="o">=</span><span class="kc">None</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Check input tensor types to be valid and confirm they are the same type for dot and batch dot ops.</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="n">msg_prefix</span> <span class="o">=</span> <span class="sa">f</span><span class="s2">&quot;For &#39;</span><span class="si">{</span><span class="n">prim_name</span><span class="si">}</span><span class="s2">&#39;, the&quot;</span> <span class="k">if</span> <span class="n">prim_name</span> <span class="k">else</span> <span class="s2">&quot;The&quot;</span>
    <span class="n">const_utils</span><span class="o">.</span><span class="n">check_type_valid</span><span class="p">(</span><span class="n">x1_type</span><span class="p">,</span> <span class="p">[</span><span class="n">mstype</span><span class="o">.</span><span class="n">float16</span><span class="p">,</span> <span class="n">mstype</span><span class="o">.</span><span class="n">float32</span><span class="p">],</span> <span class="s1">&#39;x1&#39;</span><span class="p">)</span>
    <span class="n">const_utils</span><span class="o">.</span><span class="n">check_type_valid</span><span class="p">(</span><span class="n">x2_type</span><span class="p">,</span> <span class="p">[</span><span class="n">mstype</span><span class="o">.</span><span class="n">float16</span><span class="p">,</span> <span class="n">mstype</span><span class="o">.</span><span class="n">float32</span><span class="p">],</span> <span class="s1">&#39;x2&#39;</span><span class="p">)</span>
    <span class="k">if</span> <span class="n">x1_type</span> <span class="o">!=</span> <span class="n">x2_type</span><span class="p">:</span>
        <span class="k">raise</span> <span class="ne">TypeError</span><span class="p">(</span><span class="sa">f</span><span class="s2">&quot;</span><span class="si">{</span><span class="n">msg_prefix</span><span class="si">}</span><span class="s2"> inputs must be the same type, but got &quot;</span>
                        <span class="sa">f</span><span class="s2">&quot;x1_type: </span><span class="si">{</span><span class="n">x1_type</span><span class="si">}</span><span class="s2"> and x2_type: </span><span class="si">{</span><span class="n">x2_type</span><span class="si">}</span><span class="s2">.&quot;</span><span class="p">)</span>


<span class="nd">@constexpr</span>
<span class="k">def</span> <span class="nf">_get_transpose_shape</span><span class="p">(</span><span class="n">x2_shape</span><span class="p">):</span>
    <span class="n">x2_shape_range</span> <span class="o">=</span> <span class="nb">tuple</span><span class="p">(</span><span class="nb">range</span><span class="p">(</span><span class="nb">len</span><span class="p">(</span><span class="n">x2_shape</span><span class="p">)))</span>
    <span class="n">x2_shape_transpose</span> <span class="o">=</span> <span class="n">x2_shape_range</span><span class="p">[</span><span class="o">-</span><span class="mi">2</span><span class="p">:</span><span class="o">-</span><span class="mi">1</span><span class="p">]</span> <span class="o">+</span> <span class="n">x2_shape_range</span><span class="p">[:</span><span class="o">-</span><span class="mi">2</span><span class="p">]</span> <span class="o">+</span> <span class="n">x2_shape_range</span><span class="p">[</span><span class="o">-</span><span class="mi">1</span><span class="p">:]</span>
    <span class="k">return</span> <span class="n">x2_shape_transpose</span>


<div class="viewcode-block" id="dot"><a class="viewcode-back" href="../../../../api_python/ops/mindspore.ops.dot.html#mindspore.ops.dot">[docs]</a><span class="k">def</span> <span class="nf">dot</span><span class="p">(</span><span class="n">x1</span><span class="p">,</span> <span class="n">x2</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Computation a dot product between samples in two tensors.</span>

<span class="sd">    Args:</span>
<span class="sd">        x1 (Tensor): First tensor in Dot op with datatype float16 or float32,</span>
<span class="sd">            The rank must be greater than or equal to 2.</span>
<span class="sd">        x2 (Tensor): Second tensor in Dot op with datatype float16 or float32,</span>
<span class="sd">            The rank must be greater than or equal to 2.</span>

<span class="sd">    Inputs:</span>
<span class="sd">        - **x1** (Tensor) - First tensor in Dot op with datatype float16 or float32</span>
<span class="sd">          The rank must be greater than or equal to 2.</span>
<span class="sd">        - **x2** (Tensor) - Second tensor in Dot op with datatype float16 or float32</span>
<span class="sd">          The rank must be greater than or equal to 2.</span>

<span class="sd">    Outputs:</span>
<span class="sd">        Tensor, dot product of x1 and x2.</span>

<span class="sd">    Raises:</span>
<span class="sd">        TypeError: If type of x1 and x2 are not the same.</span>
<span class="sd">        TypeError: If dtype of x1 or x2 is not float16 or float32.</span>
<span class="sd">        ValueError: If rank of x1 or x2 less than 2.</span>

<span class="sd">    Supported Platforms:</span>
<span class="sd">        ``Ascend`` ``GPU`` ``CPU``</span>

<span class="sd">    Examples:</span>
<span class="sd">        &gt;&gt;&gt; from mindspore import Tensor, ops</span>
<span class="sd">        &gt;&gt;&gt; import mindspore</span>
<span class="sd">        &gt;&gt;&gt; input_x1 = Tensor(np.ones(shape=[2, 3]), mindspore.float32)</span>
<span class="sd">        &gt;&gt;&gt; input_x2 = Tensor(np.ones(shape=[1, 3, 2]), mindspore.float32)</span>
<span class="sd">        &gt;&gt;&gt; output = ops.dot(input_x1, input_x2)</span>
<span class="sd">        &gt;&gt;&gt; print(output)</span>
<span class="sd">        [[[3. 3.]]</span>
<span class="sd">         [[3. 3.]]]</span>
<span class="sd">        &gt;&gt;&gt; print(output.shape)</span>
<span class="sd">        (2, 1, 2)</span>
<span class="sd">        &gt;&gt;&gt; input_x1 = Tensor(np.ones(shape=[1, 2, 3]), mindspore.float32)</span>
<span class="sd">        &gt;&gt;&gt; input_x2 = Tensor(np.ones(shape=[1, 3, 2]), mindspore.float32)</span>
<span class="sd">        &gt;&gt;&gt; output = ops.dot(input_x1, input_x2)</span>
<span class="sd">        &gt;&gt;&gt; print(output)</span>
<span class="sd">        [[[[3. 3.]]</span>
<span class="sd">          [[3. 3.]]]]</span>
<span class="sd">        &gt;&gt;&gt; print(output.shape)</span>
<span class="sd">        (1, 2, 1, 2)</span>
<span class="sd">        &gt;&gt;&gt; input_x1 = Tensor(np.ones(shape=[1, 2, 3]), mindspore.float32)</span>
<span class="sd">        &gt;&gt;&gt; input_x2 = Tensor(np.ones(shape=[2, 3, 2]), mindspore.float32)</span>
<span class="sd">        &gt;&gt;&gt; output = ops.dot(input_x1, input_x2)</span>
<span class="sd">        &gt;&gt;&gt; print(output)</span>
<span class="sd">        [[[[3. 3.]</span>
<span class="sd">           [3. 3.]]</span>
<span class="sd">          [[3. 3.]</span>
<span class="sd">           [3. 3.]]]]</span>
<span class="sd">        &gt;&gt;&gt; print(output.shape)</span>
<span class="sd">        (1, 2, 2, 2)</span>
<span class="sd">        &gt;&gt;&gt; input_x1 = Tensor(np.ones(shape=[3, 2, 3]), mindspore.float32)</span>
<span class="sd">        &gt;&gt;&gt; input_x2 = Tensor(np.ones(shape=[2, 1, 3, 2]), mindspore.float32)</span>
<span class="sd">        &gt;&gt;&gt; output = ops.dot(input_x1, input_x2)</span>
<span class="sd">        &gt;&gt;&gt; print(output)</span>
<span class="sd">        [[[[[3. 3.]]</span>
<span class="sd">           [[3. 3.]]]</span>
<span class="sd">          [[[3. 3.]]</span>
<span class="sd">           [[3. 3.]]]]</span>
<span class="sd">         [[[[3. 3.]]</span>
<span class="sd">           [[3. 3.]]]</span>
<span class="sd">          [[[3. 3.]]</span>
<span class="sd">           [[3. 3.]]]]</span>
<span class="sd">         [[[[3. 3.]]</span>
<span class="sd">           [[3. 3.]]]</span>
<span class="sd">          [[[3. 3.]]</span>
<span class="sd">           [[3. 3.]]]]]</span>
<span class="sd">        &gt;&gt;&gt; print(output.shape)</span>
<span class="sd">        (3, 2, 2, 1, 2)</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="n">shape_op</span> <span class="o">=</span> <span class="n">P</span><span class="o">.</span><span class="n">Shape</span><span class="p">()</span>
    <span class="n">reshape_op</span> <span class="o">=</span> <span class="n">P</span><span class="o">.</span><span class="n">Reshape</span><span class="p">()</span>
    <span class="n">transpose_op</span> <span class="o">=</span> <span class="n">P</span><span class="o">.</span><span class="n">Transpose</span><span class="p">()</span>
    <span class="n">matmul_op</span> <span class="o">=</span> <span class="n">P</span><span class="o">.</span><span class="n">MatMul</span><span class="p">(</span><span class="kc">False</span><span class="p">,</span> <span class="kc">False</span><span class="p">)</span>
    <span class="n">x1_shape</span> <span class="o">=</span> <span class="n">shape_op</span><span class="p">(</span><span class="n">x1</span><span class="p">)</span>
    <span class="n">x2_shape</span> <span class="o">=</span> <span class="n">shape_op</span><span class="p">(</span><span class="n">x2</span><span class="p">)</span>
    <span class="n">x1_type</span> <span class="o">=</span> <span class="n">F</span><span class="o">.</span><span class="n">dtype</span><span class="p">(</span><span class="n">x1</span><span class="p">)</span>
    <span class="n">x2_type</span> <span class="o">=</span> <span class="n">F</span><span class="o">.</span><span class="n">dtype</span><span class="p">(</span><span class="n">x2</span><span class="p">)</span>
    <span class="n">_typecheck_input_dot</span><span class="p">(</span><span class="n">x1_type</span><span class="p">,</span> <span class="n">x2_type</span><span class="p">,</span> <span class="s1">&#39;dot&#39;</span><span class="p">)</span>
    <span class="n">_check_invalid_input</span><span class="p">(</span><span class="n">x1_shape</span><span class="p">,</span> <span class="n">x2_shape</span><span class="p">,</span> <span class="s1">&#39;dot&#39;</span><span class="p">)</span>

    <span class="k">if</span> <span class="nb">len</span><span class="p">(</span><span class="n">x1_shape</span><span class="p">)</span> <span class="o">&gt;</span> <span class="mi">2</span> <span class="ow">or</span> <span class="nb">len</span><span class="p">(</span><span class="n">x2_shape</span><span class="p">)</span> <span class="o">&gt;</span> <span class="mi">2</span><span class="p">:</span>
        <span class="n">x2_shape_transpose</span> <span class="o">=</span> <span class="n">_get_transpose_shape</span><span class="p">(</span><span class="n">x2_shape</span><span class="p">)</span>
        <span class="n">x2_transpose</span> <span class="o">=</span> <span class="n">transpose_op</span><span class="p">(</span><span class="n">x2</span><span class="p">,</span> <span class="n">x2_shape_transpose</span><span class="p">)</span>
        <span class="n">x1_reshape</span> <span class="o">=</span> <span class="n">reshape_op</span><span class="p">(</span><span class="n">x1</span><span class="p">,</span> <span class="p">(</span><span class="o">-</span><span class="mi">1</span><span class="p">,</span> <span class="n">x1_shape</span><span class="p">[</span><span class="o">-</span><span class="mi">1</span><span class="p">]))</span>
        <span class="n">x2_reshape</span> <span class="o">=</span> <span class="n">reshape_op</span><span class="p">(</span><span class="n">x2_transpose</span><span class="p">,</span> <span class="p">(</span><span class="n">x2_shape</span><span class="p">[</span><span class="o">-</span><span class="mi">2</span><span class="p">],</span> <span class="o">-</span><span class="mi">1</span><span class="p">))</span>
        <span class="n">mul_result</span> <span class="o">=</span> <span class="n">matmul_op</span><span class="p">(</span><span class="n">x1_reshape</span><span class="p">,</span> <span class="n">x2_reshape</span><span class="p">)</span>
        <span class="k">return</span> <span class="n">reshape_op</span><span class="p">(</span><span class="n">mul_result</span><span class="p">,</span> <span class="n">x1_shape</span><span class="p">[:</span><span class="o">-</span><span class="mi">1</span><span class="p">]</span> <span class="o">+</span> <span class="n">x2_shape</span><span class="p">[:</span><span class="o">-</span><span class="mi">2</span><span class="p">]</span> <span class="o">+</span> <span class="n">x2_shape</span><span class="p">[</span><span class="o">-</span><span class="mi">1</span><span class="p">:])</span>
    <span class="k">return</span> <span class="n">matmul_op</span><span class="p">(</span><span class="n">x1</span><span class="p">,</span> <span class="n">x2</span><span class="p">)</span></div>


<span class="nd">@constexpr</span>
<span class="k">def</span> <span class="nf">_get_batch_size</span><span class="p">(</span><span class="n">x1_shape</span><span class="p">,</span> <span class="n">x2_shape</span><span class="p">,</span> <span class="n">prim_name</span><span class="o">=</span><span class="kc">None</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Get batch sizes from two inputs</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="n">msg_prefix</span> <span class="o">=</span> <span class="sa">f</span><span class="s2">&quot;For &#39;</span><span class="si">{</span><span class="n">prim_name</span><span class="si">}</span><span class="s2">&#39;, the&quot;</span> <span class="k">if</span> <span class="n">prim_name</span> <span class="k">else</span> <span class="s2">&quot;The&quot;</span>
    <span class="k">if</span> <span class="nb">len</span><span class="p">(</span><span class="n">x1_shape</span><span class="p">)</span> <span class="o">&lt;</span> <span class="mi">2</span> <span class="ow">or</span> <span class="nb">len</span><span class="p">(</span><span class="n">x2_shape</span><span class="p">)</span> <span class="o">&lt;</span> <span class="mi">2</span><span class="p">:</span>
        <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="sa">f</span><span class="s2">&quot;</span><span class="si">{</span><span class="n">msg_prefix</span><span class="si">}</span><span class="s2"> inputs x1, x2 should have &#39;dimension &gt;= 2&#39;, &quot;</span>
                         <span class="sa">f</span><span class="s2">&quot;but got &#39;len(x1_shape)&#39;: (</span><span class="si">{</span><span class="nb">len</span><span class="p">(</span><span class="n">x1_shape</span><span class="p">)</span><span class="si">}</span><span class="s2">) and &#39;len(x2_shape)&#39;: (</span><span class="si">{</span><span class="nb">len</span><span class="p">(</span><span class="n">x2_shape</span><span class="p">)</span><span class="si">}</span><span class="s2">).&quot;</span><span class="p">)</span>
    <span class="k">return</span> <span class="n">x1_shape</span><span class="p">[</span><span class="mi">0</span><span class="p">],</span> <span class="n">x2_shape</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span>


<span class="nd">@constexpr</span>
<span class="k">def</span> <span class="nf">_typecheck_input_batch_dot</span><span class="p">(</span><span class="n">x1_type</span><span class="p">,</span> <span class="n">x2_type</span><span class="p">,</span> <span class="n">prim_name</span><span class="o">=</span><span class="kc">None</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Check input tensor types to be valid and confirm they are the same type for batch dot ops.</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="n">msg_prefix</span> <span class="o">=</span> <span class="sa">f</span><span class="s2">&quot;For &#39;</span><span class="si">{</span><span class="n">prim_name</span><span class="si">}</span><span class="s2">&#39;, the&quot;</span> <span class="k">if</span> <span class="n">prim_name</span> <span class="k">else</span> <span class="s2">&quot;The&quot;</span>
    <span class="n">const_utils</span><span class="o">.</span><span class="n">check_type_valid</span><span class="p">(</span><span class="n">x1_type</span><span class="p">,</span> <span class="p">[</span><span class="n">mstype</span><span class="o">.</span><span class="n">float32</span><span class="p">],</span> <span class="s1">&#39;x1&#39;</span><span class="p">)</span>
    <span class="n">const_utils</span><span class="o">.</span><span class="n">check_type_valid</span><span class="p">(</span><span class="n">x2_type</span><span class="p">,</span> <span class="p">[</span><span class="n">mstype</span><span class="o">.</span><span class="n">float32</span><span class="p">],</span> <span class="s1">&#39;x2&#39;</span><span class="p">)</span>
    <span class="k">if</span> <span class="n">x1_type</span> <span class="o">!=</span> <span class="n">x2_type</span><span class="p">:</span>
        <span class="k">raise</span> <span class="ne">TypeError</span><span class="p">(</span><span class="sa">f</span><span class="s2">&quot;</span><span class="si">{</span><span class="n">msg_prefix</span><span class="si">}</span><span class="s2"> inputs must be the same type, but got x1_type: </span><span class="si">{</span><span class="n">x1_type</span><span class="si">}</span><span class="s2"> and &quot;</span>
                        <span class="sa">f</span><span class="s2">&quot;x2_type: </span><span class="si">{</span><span class="n">x2_type</span><span class="si">}</span><span class="s2">.&quot;</span><span class="p">)</span>


<span class="nd">@constexpr</span>
<span class="k">def</span> <span class="nf">_check_axes_for_batch_dot</span><span class="p">(</span><span class="n">x1_shape</span><span class="p">,</span> <span class="n">x2_shape</span><span class="p">,</span> <span class="n">axes</span><span class="p">,</span> <span class="n">prim_name</span><span class="o">=</span><span class="kc">None</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Check whether axes are valid and cast axes from tuple to list</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="n">msg_prefix</span> <span class="o">=</span> <span class="sa">f</span><span class="s2">&quot;For &#39;</span><span class="si">{</span><span class="n">prim_name</span><span class="si">}</span><span class="s2">&#39;, the&quot;</span> <span class="k">if</span> <span class="n">prim_name</span> <span class="k">else</span> <span class="s2">&quot;The&quot;</span>
    <span class="k">if</span> <span class="n">axes</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
        <span class="k">if</span> <span class="nb">len</span><span class="p">(</span><span class="n">x2_shape</span><span class="p">)</span> <span class="o">==</span> <span class="mi">2</span><span class="p">:</span>
            <span class="n">axes</span> <span class="o">=</span> <span class="p">[</span><span class="nb">len</span><span class="p">(</span><span class="n">x1_shape</span><span class="p">)</span> <span class="o">-</span> <span class="mi">1</span><span class="p">,</span> <span class="nb">len</span><span class="p">(</span><span class="n">x2_shape</span><span class="p">)</span> <span class="o">-</span> <span class="mi">1</span><span class="p">]</span>
        <span class="k">else</span><span class="p">:</span>
            <span class="n">axes</span> <span class="o">=</span> <span class="p">[</span><span class="nb">len</span><span class="p">(</span><span class="n">x1_shape</span><span class="p">)</span> <span class="o">-</span> <span class="mi">1</span><span class="p">,</span> <span class="nb">len</span><span class="p">(</span><span class="n">x2_shape</span><span class="p">)</span> <span class="o">-</span> <span class="mi">2</span><span class="p">]</span>

    <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">axes</span><span class="p">,</span> <span class="p">(</span><span class="nb">list</span><span class="p">,</span> <span class="nb">tuple</span><span class="p">)):</span>
        <span class="k">if</span> <span class="mi">0</span> <span class="ow">in</span> <span class="n">axes</span><span class="p">:</span>
            <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="sa">f</span><span class="s2">&quot;</span><span class="si">{</span><span class="n">msg_prefix</span><span class="si">}</span><span class="s2"> &#39;axes&#39; cannot contain 0, but got axes: </span><span class="si">{</span><span class="n">axes</span><span class="si">}</span><span class="s2">.&quot;</span><span class="p">)</span>
        <span class="k">if</span> <span class="nb">len</span><span class="p">(</span><span class="n">axes</span><span class="p">)</span> <span class="o">!=</span> <span class="mi">2</span><span class="p">:</span>
            <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="sa">f</span><span class="s2">&quot;</span><span class="si">{</span><span class="n">msg_prefix</span><span class="si">}</span><span class="s2"> length of &#39;axes&#39; must be equal to 2, but got </span><span class="si">{</span><span class="nb">len</span><span class="p">(</span><span class="n">axes</span><span class="p">)</span><span class="si">}</span><span class="s2">.&quot;</span><span class="p">)</span>
        <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">axes</span><span class="p">,</span> <span class="nb">tuple</span><span class="p">):</span>
            <span class="n">axes</span> <span class="o">=</span> <span class="nb">list</span><span class="p">(</span><span class="n">axes</span><span class="p">)</span>
        <span class="n">validator</span><span class="o">.</span><span class="n">check_value_type</span><span class="p">(</span><span class="s1">&#39;axes[0]&#39;</span><span class="p">,</span> <span class="n">axes</span><span class="p">[</span><span class="mi">0</span><span class="p">],</span> <span class="p">[</span><span class="nb">int</span><span class="p">],</span> <span class="s1">&#39;batch_dot&#39;</span><span class="p">)</span>
        <span class="n">validator</span><span class="o">.</span><span class="n">check_value_type</span><span class="p">(</span><span class="s1">&#39;axes[1]&#39;</span><span class="p">,</span> <span class="n">axes</span><span class="p">[</span><span class="mi">1</span><span class="p">],</span> <span class="p">[</span><span class="nb">int</span><span class="p">],</span> <span class="s1">&#39;batch_dot&#39;</span><span class="p">)</span>
        <span class="c1"># Reverse if axis &lt; 0</span>
        <span class="k">if</span> <span class="n">axes</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span> <span class="o">&lt;</span> <span class="mi">0</span><span class="p">:</span>
            <span class="n">axes</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span> <span class="o">+=</span> <span class="nb">len</span><span class="p">(</span><span class="n">x1_shape</span><span class="p">)</span>
        <span class="k">if</span> <span class="n">axes</span><span class="p">[</span><span class="mi">1</span><span class="p">]</span> <span class="o">&lt;</span> <span class="mi">0</span><span class="p">:</span>
            <span class="n">axes</span><span class="p">[</span><span class="mi">1</span><span class="p">]</span> <span class="o">+=</span> <span class="nb">len</span><span class="p">(</span><span class="n">x2_shape</span><span class="p">)</span>
        <span class="n">validator</span><span class="o">.</span><span class="n">check_non_negative_int</span><span class="p">(</span><span class="n">axes</span><span class="p">[</span><span class="mi">0</span><span class="p">],</span> <span class="s1">&#39;reversed axes[0]&#39;</span><span class="p">,</span> <span class="s1">&#39;batch_dot&#39;</span><span class="p">)</span>
        <span class="n">validator</span><span class="o">.</span><span class="n">check_non_negative_int</span><span class="p">(</span><span class="n">axes</span><span class="p">[</span><span class="mi">1</span><span class="p">],</span> <span class="s1">&#39;reversed axes[1]&#39;</span><span class="p">,</span> <span class="s1">&#39;batch_dot&#39;</span><span class="p">)</span>
        <span class="k">if</span> <span class="n">axes</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span> <span class="o">&gt;</span> <span class="nb">len</span><span class="p">(</span><span class="n">x1_shape</span><span class="p">)</span> <span class="ow">or</span> <span class="n">axes</span><span class="p">[</span><span class="mi">1</span><span class="p">]</span> <span class="o">&gt;</span> <span class="nb">len</span><span class="p">(</span><span class="n">x2_shape</span><span class="p">):</span>
            <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="sa">f</span><span class="s2">&quot;</span><span class="si">{</span><span class="n">msg_prefix</span><span class="si">}</span><span class="s2"> axes[0] must be less than or equal to len(x1_shape), &quot;</span>
                             <span class="sa">f</span><span class="s2">&quot;and axes[1] must be less than or equal to len(x2_shape).&quot;</span>
                             <span class="sa">f</span><span class="s2">&quot;But got &#39;axes&#39;: </span><span class="si">{</span><span class="n">axes</span><span class="si">}</span><span class="s2">, &#39;x1_shape&#39;: </span><span class="si">{</span><span class="n">x1_shape</span><span class="si">}</span><span class="s2">, &#39;x2_shape&#39;: </span><span class="si">{</span><span class="n">x2_shape</span><span class="si">}</span><span class="s2">.&quot;</span><span class="p">)</span>
    <span class="k">elif</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">axes</span><span class="p">,</span> <span class="nb">int</span><span class="p">):</span>
        <span class="k">if</span> <span class="n">axes</span> <span class="o">==</span> <span class="mi">0</span><span class="p">:</span>
            <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="sa">f</span><span class="s2">&quot;</span><span class="si">{</span><span class="n">msg_prefix</span><span class="si">}</span><span class="s2"> &#39;axes&#39; should not be equal to 0, but got </span><span class="si">{</span><span class="n">axes</span><span class="si">}</span><span class="s2">.&quot;</span><span class="p">)</span>
        <span class="k">if</span> <span class="n">axes</span> <span class="o">&lt;</span> <span class="mi">0</span><span class="p">:</span>
            <span class="n">axes</span> <span class="o">=</span> <span class="p">[</span><span class="n">axes</span> <span class="o">+</span> <span class="nb">len</span><span class="p">(</span><span class="n">x1_shape</span><span class="p">),</span> <span class="n">axes</span> <span class="o">+</span> <span class="nb">len</span><span class="p">(</span><span class="n">x2_shape</span><span class="p">)]</span>
            <span class="n">validator</span><span class="o">.</span><span class="n">check_non_negative_int</span><span class="p">(</span><span class="n">axes</span><span class="p">[</span><span class="mi">0</span><span class="p">],</span> <span class="s1">&#39;reversed axes&#39;</span><span class="p">,</span> <span class="s1">&#39;batch_dot&#39;</span><span class="p">)</span>
        <span class="k">elif</span> <span class="n">axes</span> <span class="o">&gt;</span> <span class="nb">len</span><span class="p">(</span><span class="n">x1_shape</span><span class="p">)</span> <span class="ow">or</span> <span class="n">axes</span> <span class="o">&gt;</span> <span class="nb">len</span><span class="p">(</span><span class="n">x2_shape</span><span class="p">):</span>
            <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="sa">f</span><span class="s2">&quot;</span><span class="si">{</span><span class="n">msg_prefix</span><span class="si">}</span><span class="s2"> &#39;axes&#39; cannot be greater than the length of &#39;x1_shape&#39; and &#39;x2_shape&#39;, &quot;</span>
                             <span class="sa">f</span><span class="s2">&quot;but got &#39;axes&#39;: </span><span class="si">{</span><span class="n">axes</span><span class="si">}</span><span class="s2">, &#39;x1_shape&#39;: </span><span class="si">{</span><span class="n">x1_shape</span><span class="si">}</span><span class="s2">, &#39;x2_shape&#39;: </span><span class="si">{</span><span class="n">x2_shape</span><span class="si">}</span><span class="s2">.&quot;</span><span class="p">)</span>
        <span class="k">else</span><span class="p">:</span>
            <span class="n">axes</span> <span class="o">=</span> <span class="p">[</span><span class="n">axes</span><span class="p">,</span> <span class="n">axes</span><span class="p">]</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="sa">f</span><span class="s2">&quot;</span><span class="si">{</span><span class="n">msg_prefix</span><span class="si">}</span><span class="s2"> type of &#39;axes&#39; must be one of those: int, tuple(int), list(int), &quot;</span>
                         <span class="sa">f</span><span class="s2">&quot;but got </span><span class="si">{</span><span class="nb">type</span><span class="p">(</span><span class="n">axes</span><span class="p">)</span><span class="o">.</span><span class="vm">__name__</span><span class="si">}</span><span class="s2">.&quot;</span><span class="p">)</span>
    <span class="k">return</span> <span class="n">axes</span>


<span class="nd">@constexpr</span>
<span class="k">def</span> <span class="nf">_calc_new_shape_batchdot</span><span class="p">(</span><span class="n">shape</span><span class="p">,</span> <span class="n">axes</span><span class="p">,</span> <span class="n">position</span><span class="o">=</span><span class="mi">0</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Calculate transpose and reshape parameters for input transformations,</span>
<span class="sd">    &#39;position&#39; refers to whether tensor is first or second in the op.</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="n">axis</span> <span class="o">=</span> <span class="n">axes</span><span class="p">[</span><span class="n">position</span><span class="p">]</span>
    <span class="n">contraction_axes</span> <span class="o">=</span> <span class="nb">tuple</span><span class="p">([</span><span class="n">axis</span><span class="p">])</span>
    <span class="n">prod_contraction</span> <span class="o">=</span> <span class="nb">int</span><span class="p">(</span><span class="n">np</span><span class="o">.</span><span class="n">prod</span><span class="p">([</span><span class="n">shape</span><span class="p">[</span><span class="n">i</span><span class="p">]</span> <span class="k">for</span> <span class="n">i</span> <span class="ow">in</span> <span class="n">contraction_axes</span><span class="p">]))</span>
    <span class="n">free_axes</span> <span class="o">=</span> <span class="nb">tuple</span><span class="p">(</span><span class="n">i</span> <span class="k">for</span> <span class="n">i</span> <span class="ow">in</span> <span class="nb">range</span><span class="p">(</span><span class="mi">1</span><span class="p">,</span> <span class="nb">len</span><span class="p">(</span><span class="n">shape</span><span class="p">))</span> <span class="k">if</span> <span class="n">i</span> <span class="ow">not</span> <span class="ow">in</span> <span class="n">contraction_axes</span><span class="p">)</span>
    <span class="n">free_dims</span> <span class="o">=</span> <span class="nb">tuple</span><span class="p">(</span><span class="n">shape</span><span class="p">[</span><span class="n">i</span><span class="p">]</span> <span class="k">for</span> <span class="n">i</span> <span class="ow">in</span> <span class="n">free_axes</span><span class="p">)</span>
    <span class="n">prod_free</span> <span class="o">=</span> <span class="nb">int</span><span class="p">(</span><span class="n">np</span><span class="o">.</span><span class="n">prod</span><span class="p">(</span><span class="n">free_dims</span><span class="p">))</span>

    <span class="n">transpose_perm</span> <span class="o">=</span> <span class="n">contraction_axes</span> <span class="o">+</span> <span class="n">free_axes</span> <span class="k">if</span> <span class="n">position</span> <span class="k">else</span> <span class="n">free_axes</span> <span class="o">+</span> <span class="n">contraction_axes</span>
    <span class="n">transpose_perm</span> <span class="o">=</span> <span class="nb">tuple</span><span class="p">([</span><span class="mi">0</span><span class="p">])</span> <span class="o">+</span> <span class="n">transpose_perm</span>
    <span class="n">new_shape</span> <span class="o">=</span> <span class="p">(</span><span class="n">prod_contraction</span><span class="p">,</span> <span class="n">prod_free</span><span class="p">)</span> <span class="k">if</span> <span class="n">position</span> <span class="k">else</span> <span class="p">(</span><span class="n">prod_free</span><span class="p">,</span> <span class="n">prod_contraction</span><span class="p">)</span>
    <span class="n">new_shape</span> <span class="o">=</span> <span class="nb">tuple</span><span class="p">([</span><span class="n">shape</span><span class="p">[</span><span class="mi">0</span><span class="p">]])</span> <span class="o">+</span> <span class="n">new_shape</span>
    <span class="k">return</span> <span class="n">new_shape</span><span class="p">,</span> <span class="n">transpose_perm</span><span class="p">,</span> <span class="n">free_dims</span>


<span class="nd">@constexpr</span>
<span class="k">def</span> <span class="nf">_check_batch_size</span><span class="p">(</span><span class="n">x1_batch_size</span><span class="p">,</span> <span class="n">x2_batch_size</span><span class="p">,</span> <span class="n">prim_name</span><span class="o">=</span><span class="kc">None</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Check whether batch size of two inputs are the same</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="n">msg_prefix</span> <span class="o">=</span> <span class="sa">f</span><span class="s2">&quot;For &#39;</span><span class="si">{</span><span class="n">prim_name</span><span class="si">}</span><span class="s2">&#39;, the&quot;</span> <span class="k">if</span> <span class="n">prim_name</span> <span class="k">else</span> <span class="s2">&quot;The&quot;</span>
    <span class="k">if</span> <span class="n">x1_batch_size</span> <span class="o">!=</span> <span class="n">x2_batch_size</span><span class="p">:</span>
        <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="sa">f</span><span class="s2">&quot;</span><span class="si">{</span><span class="n">msg_prefix</span><span class="si">}</span><span class="s2"> inputs &#39;x1&#39;, &#39;x2&#39; should have the same batch sizes, but got &quot;</span>
                         <span class="sa">f</span><span class="s2">&quot;&#39;x1_batch_size&#39;: </span><span class="si">{</span><span class="n">x1_batch_size</span><span class="si">}</span><span class="s2"> and &#39;x2_batch_size&#39;: </span><span class="si">{</span><span class="n">x2_batch_size</span><span class="si">}</span><span class="s2">.&quot;</span><span class="p">)</span>


<span class="nd">@constexpr</span>
<span class="k">def</span> <span class="nf">_get_output_shape</span><span class="p">(</span><span class="n">batch_size</span><span class="p">,</span> <span class="n">x1_ret</span><span class="p">,</span> <span class="n">x2_ret</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Compute output shape for batch dot</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="n">output_shape</span> <span class="o">=</span> <span class="nb">tuple</span><span class="p">([</span><span class="n">batch_size</span><span class="p">])</span> <span class="o">+</span> <span class="n">x1_ret</span> <span class="o">+</span> <span class="n">x2_ret</span>
    <span class="k">return</span> <span class="n">output_shape</span>


<div class="viewcode-block" id="batch_dot"><a class="viewcode-back" href="../../../../api_python/ops/mindspore.ops.batch_dot.html#mindspore.ops.batch_dot">[docs]</a><span class="k">def</span> <span class="nf">batch_dot</span><span class="p">(</span><span class="n">x1</span><span class="p">,</span> <span class="n">x2</span><span class="p">,</span> <span class="n">axes</span><span class="o">=</span><span class="kc">None</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Computation of batch dot product between samples in two tensors containing batch dims.</span>

<span class="sd">    .. math::</span>
<span class="sd">        output = x1[batch, :] * x2[batch, :]</span>

<span class="sd">    Args:</span>
<span class="sd">        x1 (Tensor): First tensor in Batch Dot op with datatype float32 and the rank of `x1` must be greater</span>
<span class="sd">          than or equal to 2.</span>
<span class="sd">        x2 (Tensor): Second tensor in Batch Dot op with datatype float32. The datatype of `x2` should</span>
<span class="sd">          be same as `x1` and the rank of `x2` must be greater than or equal to 2.</span>
<span class="sd">        axes (Union[int, tuple(int), list(int)]): Single value or tuple/list of length 2 with dimensions</span>
<span class="sd">          specified for `a` and `b` each. If single value `N` passed, automatically picks up last N dims from</span>
<span class="sd">          `a` input shape and last N dimensions from `b` input shape in order as axes for each respectively.</span>
<span class="sd">          Default: None.</span>

<span class="sd">    Outputs:</span>
<span class="sd">        Tensor, batch dot product of `x1` and `x2`. For example: The Shape of output</span>
<span class="sd">        for input `x1` shapes (batch, d1, axes, d2) and `x2` shapes (batch, d3, axes, d4) is (batch, d1, d2, d3, d4),</span>
<span class="sd">        where d1 and d2 means any number.</span>

<span class="sd">    Raises:</span>
<span class="sd">        TypeError: If type of x1 and x2 are not the same.</span>
<span class="sd">        TypeError: If dtype of x1 or x2 is not float32.</span>
<span class="sd">        ValueError: If rank of x1 or x2 less than 2.</span>
<span class="sd">        ValueError: If batch dim used in axes.</span>
<span class="sd">        ValueError: If len(axes) less than 2.</span>
<span class="sd">        ValueError: If axes is not one of those: None, int, (int, int).</span>
<span class="sd">        ValueError: If axes reversed from negative int is too low for dimensions of input arrays.</span>
<span class="sd">        ValueError: If axes value is too high for dimensions of input arrays.</span>
<span class="sd">        ValueError: If batch size of x1 and x2 are not the same.</span>

<span class="sd">    Supported Platforms:</span>
<span class="sd">        ``Ascend`` ``GPU`` ``CPU``</span>

<span class="sd">    Examples:</span>
<span class="sd">        &gt;&gt;&gt; from mindspore import Tensor, ops</span>
<span class="sd">        &gt;&gt;&gt; import numpy as np</span>
<span class="sd">        &gt;&gt;&gt; x1 = Tensor(np.ones(shape=[2, 2, 3]), mindspore.float32)</span>
<span class="sd">        &gt;&gt;&gt; x2 = Tensor(np.ones(shape=[2, 3, 2]), mindspore.float32)</span>
<span class="sd">        &gt;&gt;&gt; axes = (-1, -2)</span>
<span class="sd">        &gt;&gt;&gt; output = ops.batch_dot(x1, x2, axes)</span>
<span class="sd">        &gt;&gt;&gt; print(output)</span>
<span class="sd">        [[[3. 3.]</span>
<span class="sd">          [3. 3.]]</span>
<span class="sd">         [[3. 3.]</span>
<span class="sd">          [3. 3.]]]</span>
<span class="sd">        &gt;&gt;&gt; x1 = Tensor(np.ones(shape=[2, 2]), mindspore.float32)</span>
<span class="sd">        &gt;&gt;&gt; x2 = Tensor(np.ones(shape=[2, 3, 2]), mindspore.float32)</span>
<span class="sd">        &gt;&gt;&gt; axes = (1, 2)</span>
<span class="sd">        &gt;&gt;&gt; output = ops.batch_dot(x1, x2, axes)</span>
<span class="sd">        &gt;&gt;&gt; print(output)</span>
<span class="sd">        [[2. 2. 2.]</span>
<span class="sd">         [2. 2. 2.]]</span>
<span class="sd">        &gt;&gt;&gt; print(output.shape)</span>
<span class="sd">        (2, 3)</span>
<span class="sd">        &gt;&gt;&gt; x1 = Tensor(np.ones(shape=[6, 2, 3, 4]), mindspore.float32)</span>
<span class="sd">        &gt;&gt;&gt; x2 = Tensor(np.ones(shape=[6, 5, 4, 8]), mindspore.float32)</span>
<span class="sd">        &gt;&gt;&gt; output = ops.batch_dot(x1, x2)</span>
<span class="sd">        &gt;&gt;&gt; print(output.shape)</span>
<span class="sd">        (6, 2, 3, 5, 8)</span>
<span class="sd">        &gt;&gt;&gt; x1 = Tensor(np.ones(shape=[2, 2, 4]), mindspore.float32)</span>
<span class="sd">        &gt;&gt;&gt; x2 = Tensor(np.ones(shape=[2, 5, 4, 5]), mindspore.float32)</span>
<span class="sd">        &gt;&gt;&gt; output = ops.batch_dot(x1, x2)</span>
<span class="sd">        &gt;&gt;&gt; print(output.shape)</span>
<span class="sd">        (2, 2, 5, 5)</span>

<span class="sd">    &quot;&quot;&quot;</span>

    <span class="n">transpose_op</span> <span class="o">=</span> <span class="n">P</span><span class="o">.</span><span class="n">Transpose</span><span class="p">()</span>
    <span class="n">batch_matmul_op</span> <span class="o">=</span> <span class="n">P</span><span class="o">.</span><span class="n">BatchMatMul</span><span class="p">()</span>
    <span class="n">squeeze_one_op</span> <span class="o">=</span> <span class="n">P</span><span class="o">.</span><span class="n">Squeeze</span><span class="p">(</span><span class="mi">1</span><span class="p">)</span>
    <span class="n">squeeze_minus_one_op</span> <span class="o">=</span> <span class="n">P</span><span class="o">.</span><span class="n">Squeeze</span><span class="p">(</span><span class="o">-</span><span class="mi">1</span><span class="p">)</span>
    <span class="c1"># input validity checks</span>
    <span class="n">x1_shape</span> <span class="o">=</span> <span class="n">F</span><span class="o">.</span><span class="n">shape</span><span class="p">(</span><span class="n">x1</span><span class="p">)</span>
    <span class="n">x2_shape</span> <span class="o">=</span> <span class="n">F</span><span class="o">.</span><span class="n">shape</span><span class="p">(</span><span class="n">x2</span><span class="p">)</span>
    <span class="n">x1_dim_num</span> <span class="o">=</span> <span class="nb">len</span><span class="p">(</span><span class="n">x1_shape</span><span class="p">)</span>
    <span class="n">x2_dim_num</span> <span class="o">=</span> <span class="nb">len</span><span class="p">(</span><span class="n">x2_shape</span><span class="p">)</span>
    <span class="n">x1_type</span> <span class="o">=</span> <span class="n">F</span><span class="o">.</span><span class="n">dtype</span><span class="p">(</span><span class="n">x1</span><span class="p">)</span>
    <span class="n">x2_type</span> <span class="o">=</span> <span class="n">F</span><span class="o">.</span><span class="n">dtype</span><span class="p">(</span><span class="n">x2</span><span class="p">)</span>

    <span class="n">x1_batch_size</span><span class="p">,</span> <span class="n">x2_batch_size</span> <span class="o">=</span> <span class="n">_get_batch_size</span><span class="p">(</span><span class="n">x1_shape</span><span class="p">,</span> <span class="n">x2_shape</span><span class="p">,</span> <span class="s1">&#39;batch_dot&#39;</span><span class="p">)</span>

    <span class="n">_typecheck_input_batch_dot</span><span class="p">(</span><span class="n">x1_type</span><span class="p">,</span> <span class="n">x2_type</span><span class="p">,</span> <span class="s1">&#39;batch_dot&#39;</span><span class="p">)</span>
    <span class="n">_check_batch_size</span><span class="p">(</span><span class="n">x1_batch_size</span><span class="p">,</span> <span class="n">x2_batch_size</span><span class="p">,</span> <span class="s1">&#39;batch_dot&#39;</span><span class="p">)</span>
    <span class="n">axes</span> <span class="o">=</span> <span class="n">_check_axes_for_batch_dot</span><span class="p">(</span><span class="n">x1_shape</span><span class="p">,</span> <span class="n">x2_shape</span><span class="p">,</span> <span class="n">axes</span><span class="p">,</span> <span class="s1">&#39;batch_dot&#39;</span><span class="p">)</span>

    <span class="k">if</span> <span class="n">x1_dim_num</span> <span class="o">==</span> <span class="mi">2</span><span class="p">:</span>
        <span class="n">x1</span> <span class="o">=</span> <span class="n">F</span><span class="o">.</span><span class="n">expand_dims</span><span class="p">(</span><span class="n">x1</span><span class="p">,</span> <span class="mi">1</span><span class="p">)</span>
        <span class="n">axes</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span> <span class="o">+=</span> <span class="mi">1</span>
    <span class="k">if</span> <span class="n">x2_dim_num</span> <span class="o">==</span> <span class="mi">2</span><span class="p">:</span>
        <span class="n">x2</span> <span class="o">=</span> <span class="n">F</span><span class="o">.</span><span class="n">expand_dims</span><span class="p">(</span><span class="n">x2</span><span class="p">,</span> <span class="mi">2</span><span class="p">)</span>

    <span class="n">x1_shape</span> <span class="o">=</span> <span class="n">F</span><span class="o">.</span><span class="n">shape</span><span class="p">(</span><span class="n">x1</span><span class="p">)</span>
    <span class="n">x2_shape</span> <span class="o">=</span> <span class="n">F</span><span class="o">.</span><span class="n">shape</span><span class="p">(</span><span class="n">x2</span><span class="p">)</span>

    <span class="n">x1_reshape_fwd</span><span class="p">,</span> <span class="n">x1_transpose_fwd</span><span class="p">,</span> <span class="n">x1_ret</span> <span class="o">=</span> <span class="n">_calc_new_shape_batchdot</span><span class="p">(</span><span class="n">x1_shape</span><span class="p">,</span> <span class="n">axes</span><span class="p">,</span> <span class="mi">0</span><span class="p">)</span>
    <span class="n">x2_reshape_fwd</span><span class="p">,</span> <span class="n">x2_transpose_fwd</span><span class="p">,</span> <span class="n">x2_ret</span> <span class="o">=</span> <span class="n">_calc_new_shape_batchdot</span><span class="p">(</span><span class="n">x2_shape</span><span class="p">,</span> <span class="n">axes</span><span class="p">,</span> <span class="mi">1</span><span class="p">)</span>
    <span class="n">output_shape</span> <span class="o">=</span> <span class="n">_get_output_shape</span><span class="p">(</span><span class="n">x1_batch_size</span><span class="p">,</span> <span class="n">x1_ret</span><span class="p">,</span> <span class="n">x2_ret</span><span class="p">)</span>

    <span class="n">x1_transposed</span> <span class="o">=</span> <span class="n">transpose_op</span><span class="p">(</span><span class="n">x1</span><span class="p">,</span> <span class="n">x1_transpose_fwd</span><span class="p">)</span>
    <span class="n">x2_transposed</span> <span class="o">=</span> <span class="n">transpose_op</span><span class="p">(</span><span class="n">x2</span><span class="p">,</span> <span class="n">x2_transpose_fwd</span><span class="p">)</span>
    <span class="n">x1_reshaped</span> <span class="o">=</span> <span class="n">F</span><span class="o">.</span><span class="n">reshape</span><span class="p">(</span><span class="n">x1_transposed</span><span class="p">,</span> <span class="n">x1_reshape_fwd</span><span class="p">)</span>
    <span class="n">x2_reshaped</span> <span class="o">=</span> <span class="n">F</span><span class="o">.</span><span class="n">reshape</span><span class="p">(</span><span class="n">x2_transposed</span><span class="p">,</span> <span class="n">x2_reshape_fwd</span><span class="p">)</span>

    <span class="c1"># Batch matmal op part</span>
    <span class="n">mul_result</span> <span class="o">=</span> <span class="n">batch_matmul_op</span><span class="p">(</span><span class="n">x1_reshaped</span><span class="p">,</span> <span class="n">x2_reshaped</span><span class="p">)</span>

    <span class="n">final_result</span> <span class="o">=</span> <span class="n">F</span><span class="o">.</span><span class="n">reshape</span><span class="p">(</span><span class="n">mul_result</span><span class="p">,</span> <span class="n">output_shape</span><span class="p">)</span>

    <span class="c1"># if the original dims are expanded, restore them from 3 to 2</span>
    <span class="k">if</span> <span class="n">x1_dim_num</span> <span class="o">==</span> <span class="mi">2</span><span class="p">:</span>
        <span class="n">final_result</span> <span class="o">=</span> <span class="n">squeeze_one_op</span><span class="p">(</span><span class="n">final_result</span><span class="p">)</span>
    <span class="k">elif</span> <span class="n">x2_dim_num</span> <span class="o">==</span> <span class="mi">2</span><span class="p">:</span>
        <span class="n">final_result</span> <span class="o">=</span> <span class="n">squeeze_minus_one_op</span><span class="p">(</span><span class="n">final_result</span><span class="p">)</span>

    <span class="k">return</span> <span class="n">final_result</span></div>


<span class="nd">@constexpr</span>
<span class="k">def</span> <span class="nf">_check_same_type</span><span class="p">(</span><span class="n">dtype1</span><span class="p">,</span> <span class="n">dtype2</span><span class="p">):</span>
    <span class="k">return</span> <span class="n">dtype1</span> <span class="o">==</span> <span class="n">dtype2</span>


<span class="nd">@constexpr</span>
<span class="k">def</span> <span class="nf">_max</span><span class="p">(</span><span class="o">*</span><span class="n">args</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;Returns the maximum value.&quot;&quot;&quot;</span>
    <span class="k">return</span> <span class="nb">max</span><span class="p">(</span><span class="o">*</span><span class="n">args</span><span class="p">)</span>


<span class="nd">@constexpr</span>
<span class="k">def</span> <span class="nf">_min</span><span class="p">(</span><span class="o">*</span><span class="n">args</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;Returns the minimum value.&quot;&quot;&quot;</span>
    <span class="k">return</span> <span class="nb">min</span><span class="p">(</span><span class="o">*</span><span class="n">args</span><span class="p">)</span>


<span class="nd">@constexpr</span>
<span class="k">def</span> <span class="nf">_infer_shape_rem</span><span class="p">(</span><span class="n">shape1</span><span class="p">,</span> <span class="n">shape2</span><span class="p">,</span> <span class="n">ndim1</span><span class="p">,</span> <span class="n">ndim2</span><span class="p">,</span> <span class="n">transpose_b</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;Infers the shape of the last two dimensions after performing matmul.&quot;&quot;&quot;</span>
    <span class="n">shape_rem</span> <span class="o">=</span> <span class="p">[]</span>
    <span class="k">if</span> <span class="n">ndim1</span> <span class="o">&gt;=</span> <span class="mi">2</span><span class="p">:</span>
        <span class="n">shape_rem</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">shape1</span><span class="p">[</span><span class="o">-</span><span class="mi">2</span><span class="p">])</span>
    <span class="k">if</span> <span class="n">transpose_b</span><span class="p">:</span>
        <span class="k">if</span> <span class="n">ndim2</span> <span class="o">&gt;=</span> <span class="mi">2</span><span class="p">:</span>
            <span class="n">shape_rem</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">shape2</span><span class="p">[</span><span class="o">-</span><span class="mi">2</span><span class="p">])</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="k">if</span> <span class="n">ndim1</span> <span class="o">&gt;=</span> <span class="mi">1</span><span class="p">:</span>
            <span class="n">shape_rem</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">shape2</span><span class="p">[</span><span class="o">-</span><span class="mi">1</span><span class="p">])</span>
    <span class="k">return</span> <span class="nb">tuple</span><span class="p">(</span><span class="n">shape_rem</span><span class="p">)</span>


<span class="nd">@constexpr</span>
<span class="k">def</span> <span class="nf">_check_matmul_shapes</span><span class="p">(</span><span class="n">shape1</span><span class="p">,</span> <span class="n">shape2</span><span class="p">,</span> <span class="n">prim_name</span><span class="o">=</span><span class="kc">None</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;Checks shape1 and shape2 are valid to perform matmul, and returns output shape after broadcasting.&quot;&quot;&quot;</span>
    <span class="n">msg_prefix</span> <span class="o">=</span> <span class="sa">f</span><span class="s2">&quot;For &#39;</span><span class="si">{</span><span class="n">prim_name</span><span class="si">}</span><span class="s2">&#39;, the&quot;</span> <span class="k">if</span> <span class="n">prim_name</span> <span class="k">else</span> <span class="s2">&quot;The&quot;</span>
    <span class="n">ndim1</span><span class="p">,</span> <span class="n">ndim2</span> <span class="o">=</span> <span class="nb">len</span><span class="p">(</span><span class="n">shape1</span><span class="p">),</span> <span class="nb">len</span><span class="p">(</span><span class="n">shape2</span><span class="p">)</span>
    <span class="k">if</span> <span class="n">ndim1</span> <span class="o">&lt;</span> <span class="mi">1</span> <span class="ow">or</span> <span class="n">ndim2</span> <span class="o">&lt;</span> <span class="mi">1</span><span class="p">:</span>
        <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="sa">f</span><span class="s2">&quot;</span><span class="si">{</span><span class="n">msg_prefix</span><span class="si">}</span><span class="s2"> dimension of input operands must be at least 1, but got &quot;</span>
                         <span class="sa">f</span><span class="s2">&quot;the length of shape1: </span><span class="si">{</span><span class="n">ndim1</span><span class="si">}</span><span class="s2">, the length of shape2: </span><span class="si">{</span><span class="n">ndim2</span><span class="si">}</span><span class="s2">.&quot;</span><span class="p">)</span>
    <span class="k">if</span> <span class="n">ndim2</span> <span class="o">&gt;=</span> <span class="mi">2</span> <span class="ow">and</span> <span class="n">shape1</span><span class="p">[</span><span class="o">-</span><span class="mi">1</span><span class="p">]</span> <span class="o">!=</span> <span class="n">shape2</span><span class="p">[</span><span class="o">-</span><span class="mi">2</span><span class="p">]:</span>
        <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="sa">f</span><span class="s2">&quot;</span><span class="si">{</span><span class="n">msg_prefix</span><span class="si">}</span><span class="s2"> shape1[-1] should be equal to shape2[-2] when the length of shape2 &quot;</span>
                         <span class="sa">f</span><span class="s2">&quot;is greater than or equal to 2, but got shape1[-1]: </span><span class="si">{</span><span class="n">shape1</span><span class="p">[</span><span class="o">-</span><span class="mi">1</span><span class="p">]</span><span class="si">}</span><span class="s2">, &quot;</span>
                         <span class="sa">f</span><span class="s2">&quot;shape2[-2]: </span><span class="si">{</span><span class="n">shape2</span><span class="p">[</span><span class="o">-</span><span class="mi">2</span><span class="p">]</span><span class="si">}</span><span class="s2">.&quot;</span><span class="p">)</span>
    <span class="n">shape_out</span> <span class="o">=</span> <span class="n">deque</span><span class="p">()</span>
    <span class="k">for</span> <span class="n">items</span> <span class="ow">in</span> <span class="n">zip_longest</span><span class="p">(</span><span class="nb">reversed</span><span class="p">(</span><span class="n">shape1</span><span class="p">[:</span><span class="o">-</span><span class="mi">2</span><span class="p">]),</span> <span class="nb">reversed</span><span class="p">(</span><span class="n">shape2</span><span class="p">[:</span><span class="o">-</span><span class="mi">2</span><span class="p">]),</span> <span class="n">fillvalue</span><span class="o">=</span><span class="mi">1</span><span class="p">):</span>
        <span class="n">max_size</span> <span class="o">=</span> <span class="nb">max</span><span class="p">(</span><span class="n">items</span><span class="p">)</span>
        <span class="k">if</span> <span class="nb">any</span><span class="p">(</span><span class="n">item</span> <span class="ow">not</span> <span class="ow">in</span> <span class="p">(</span><span class="mi">1</span><span class="p">,</span> <span class="n">max_size</span><span class="p">)</span> <span class="k">for</span> <span class="n">item</span> <span class="ow">in</span> <span class="n">items</span><span class="p">):</span>
            <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="sa">f</span><span class="s2">&quot;</span><span class="si">{</span><span class="n">msg_prefix</span><span class="si">}</span><span class="s2"> operands could not be broadcast together with shape1 </span><span class="si">{</span><span class="n">shape1</span><span class="si">}</span><span class="s2"> and &quot;</span>
                             <span class="sa">f</span><span class="s2">&quot;shape2 </span><span class="si">{</span><span class="n">shape2</span><span class="si">}</span><span class="s2">.&quot;</span><span class="p">)</span>
        <span class="n">shape_out</span><span class="o">.</span><span class="n">appendleft</span><span class="p">(</span><span class="n">max_size</span><span class="p">)</span>
    <span class="k">return</span> <span class="nb">tuple</span><span class="p">(</span><span class="n">shape_out</span><span class="p">)</span>


<span class="nd">@constexpr</span>
<span class="k">def</span> <span class="nf">_tile_size</span><span class="p">(</span><span class="n">shape</span><span class="p">,</span> <span class="n">out_shape</span><span class="p">,</span> <span class="n">ndim</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;Returns tile_size such that shape*tile_size = out_shape&quot;&quot;&quot;</span>
    <span class="n">size</span> <span class="o">=</span> <span class="p">[</span><span class="mi">1</span><span class="p">]</span> <span class="o">*</span> <span class="n">ndim</span>
    <span class="k">for</span> <span class="n">idx</span><span class="p">,</span> <span class="p">(</span><span class="n">i</span><span class="p">,</span> <span class="n">j</span><span class="p">)</span> <span class="ow">in</span> <span class="nb">enumerate</span><span class="p">(</span><span class="nb">zip</span><span class="p">(</span><span class="n">shape</span><span class="p">,</span> <span class="n">out_shape</span><span class="p">)):</span>
        <span class="k">if</span> <span class="n">i</span> <span class="o">!=</span> <span class="n">j</span><span class="p">:</span>
            <span class="n">size</span><span class="p">[</span><span class="n">idx</span><span class="p">]</span> <span class="o">=</span> <span class="n">j</span>
    <span class="k">return</span> <span class="nb">tuple</span><span class="p">(</span><span class="n">size</span><span class="p">)</span>


<span class="nd">@constexpr</span>
<span class="k">def</span> <span class="nf">_check_need_broadcast</span><span class="p">(</span><span class="n">shape1</span><span class="p">,</span> <span class="n">shape2</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;Returns True if broadcast is necessary for batchmatmul.&quot;&quot;&quot;</span>
    <span class="k">return</span> <span class="n">shape1</span><span class="p">[:</span><span class="o">-</span><span class="mi">2</span><span class="p">]</span> <span class="o">!=</span> <span class="n">shape2</span><span class="p">[:</span><span class="o">-</span><span class="mi">2</span><span class="p">]</span>


<span class="k">def</span> <span class="nf">_expand</span><span class="p">(</span><span class="n">x</span><span class="p">,</span> <span class="n">ndim</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;Expand x to ndim from axis, which can be 0 or -1.&quot;&quot;&quot;</span>
    <span class="k">while</span> <span class="n">F</span><span class="o">.</span><span class="n">rank</span><span class="p">(</span><span class="n">x</span><span class="p">)</span> <span class="o">&lt;</span> <span class="n">ndim</span><span class="p">:</span>
        <span class="n">x</span> <span class="o">=</span> <span class="n">F</span><span class="o">.</span><span class="n">expand_dims</span><span class="p">(</span><span class="n">x</span><span class="p">,</span> <span class="mi">0</span><span class="p">)</span>
    <span class="k">return</span> <span class="n">x</span>


<span class="k">def</span> <span class="nf">_broadcast_to</span><span class="p">(</span><span class="n">x</span><span class="p">,</span> <span class="n">shape_cur</span><span class="p">,</span> <span class="n">shape_to</span><span class="p">,</span> <span class="n">ndim_to</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;Broadcasts x from shape_cur to shape_to.&quot;&quot;&quot;</span>
    <span class="n">size</span> <span class="o">=</span> <span class="n">_tile_size</span><span class="p">(</span><span class="n">shape_cur</span><span class="p">,</span> <span class="n">shape_to</span><span class="p">,</span> <span class="n">ndim_to</span><span class="p">)</span>
    <span class="k">return</span> <span class="n">F</span><span class="o">.</span><span class="n">tile</span><span class="p">(</span><span class="n">x</span><span class="p">,</span> <span class="n">size</span><span class="p">)</span>


<div class="viewcode-block" id="matmul"><a class="viewcode-back" href="../../../../api_python/ops/mindspore.ops.matmul.html#mindspore.ops.matmul">[docs]</a><span class="k">def</span> <span class="nf">matmul</span><span class="p">(</span><span class="n">x1</span><span class="p">,</span> <span class="n">x2</span><span class="p">,</span> <span class="n">dtype</span><span class="o">=</span><span class="kc">None</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Returns the matrix product of two arrays.</span>

<span class="sd">    Note:</span>
<span class="sd">        Numpy arguments `out`, `casting`, `order`, `subok`, `signature`, and `extobj` are</span>
<span class="sd">        not supported.</span>
<span class="sd">        On GPU, the supported dtypes are np.float16 and np.float32.</span>
<span class="sd">        On CPU, the supported dtypes are np.float16 and np.float32.</span>

<span class="sd">    Args:</span>
<span class="sd">        x1 (Tensor): Input tensor, scalar not allowed.</span>
<span class="sd">          The last dimension of `x1` must be the same size as the second last dimension of `x2`.</span>
<span class="sd">          And the shape of x1 and x2 could be broadcast.</span>
<span class="sd">        x2 (Tensor): Input tensor, scalar not allowed.</span>
<span class="sd">          The last dimension of `x1` must be the same size as the second last dimension of `x2`.</span>
<span class="sd">          And the shape of x1 and x2 could be broadcast.</span>
<span class="sd">        dtype (:class:`mindspore.dtype`, optional): defaults to None. Overrides the dtype of the</span>
<span class="sd">            output Tensor.</span>

<span class="sd">    Returns:</span>
<span class="sd">        Tensor or scalar, the matrix product of the inputs. This is a scalar only</span>
<span class="sd">        when both `x1`, `x2` are 1-d vectors.</span>

<span class="sd">    Raises:</span>
<span class="sd">        ValueError: If the last dimension of `x1` is not the same size as the</span>
<span class="sd">            second-to-last dimension of `x2`, or if a scalar value is passed in.</span>
<span class="sd">        ValueError: If the shape of `x1` and `x2` could not broadcast together.</span>

<span class="sd">    Supported Platforms:</span>
<span class="sd">        ``Ascend`` ``GPU`` ``CPU``</span>

<span class="sd">    Examples:</span>
<span class="sd">        &gt;&gt;&gt; from mindspore import Tensor, ops</span>
<span class="sd">        &gt;&gt;&gt; import mindspore</span>
<span class="sd">        &gt;&gt;&gt; # case 1 : Reasonable application of broadcast mechanism</span>
<span class="sd">        &gt;&gt;&gt; x1 = Tensor(np.arange(2*3*4).reshape(2, 3, 4), mindspore.float32)</span>
<span class="sd">        &gt;&gt;&gt; x2 = Tensor(np.arange(4*5).reshape(4, 5), mindspore.float32)</span>
<span class="sd">        &gt;&gt;&gt; output = ops.matmul(x1, x2)</span>
<span class="sd">        &gt;&gt;&gt; print(output)</span>
<span class="sd">        [[[  70.   76.   82.   88.   94.]</span>
<span class="sd">        [ 190.  212.  234.  256.  278.]</span>
<span class="sd">        [ 310.  348.  386.  424.  462.]]</span>
<span class="sd">        [[ 430.  484.  538.  592.  646.]</span>
<span class="sd">        [ 550.  620.  690.  760.  830.]</span>
<span class="sd">        [ 670.  756.  842.  928. 1014.]]]</span>
<span class="sd">        &gt;&gt;&gt; print(output.shape)</span>
<span class="sd">        (2, 3, 5)</span>
<span class="sd">        &gt;&gt;&gt; # case 2 : the rank of `x1` is 1</span>
<span class="sd">        &gt;&gt;&gt; x1 = Tensor(np.ones([1, 2]), mindspore.float32)</span>
<span class="sd">        &gt;&gt;&gt; x2 = Tensor(np.ones([2,]), mindspore.float32)</span>
<span class="sd">        &gt;&gt;&gt; output = ops.matmul(x1, x2)</span>
<span class="sd">        &gt;&gt;&gt; print(output)</span>
<span class="sd">        [2.]</span>
<span class="sd">        &gt;&gt;&gt; print(output.shape)</span>
<span class="sd">        (1,)</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="c1"># performs type promotion</span>
    <span class="n">dtype1</span> <span class="o">=</span> <span class="n">F</span><span class="o">.</span><span class="n">dtype</span><span class="p">(</span><span class="n">x1</span><span class="p">)</span>
    <span class="n">dtype2</span> <span class="o">=</span> <span class="n">F</span><span class="o">.</span><span class="n">dtype</span><span class="p">(</span><span class="n">x2</span><span class="p">)</span>
    <span class="k">if</span> <span class="ow">not</span> <span class="n">_check_same_type</span><span class="p">(</span><span class="n">dtype1</span><span class="p">,</span> <span class="n">dtype2</span><span class="p">):</span>
        <span class="n">x1</span> <span class="o">=</span> <span class="n">x1</span><span class="o">.</span><span class="n">astype</span><span class="p">(</span><span class="n">mstype</span><span class="o">.</span><span class="n">float32</span><span class="p">)</span>
        <span class="n">x2</span> <span class="o">=</span> <span class="n">x2</span><span class="o">.</span><span class="n">astype</span><span class="p">(</span><span class="n">mstype</span><span class="o">.</span><span class="n">float32</span><span class="p">)</span>

    <span class="n">ndim1_orig</span><span class="p">,</span> <span class="n">ndim2_orig</span> <span class="o">=</span> <span class="n">F</span><span class="o">.</span><span class="n">rank</span><span class="p">(</span><span class="n">x1</span><span class="p">),</span> <span class="n">F</span><span class="o">.</span><span class="n">rank</span><span class="p">(</span><span class="n">x2</span><span class="p">)</span>
    <span class="n">shape1_orig</span><span class="p">,</span> <span class="n">shape2_orig</span> <span class="o">=</span> <span class="n">F</span><span class="o">.</span><span class="n">shape</span><span class="p">(</span><span class="n">x1</span><span class="p">),</span> <span class="n">F</span><span class="o">.</span><span class="n">shape</span><span class="p">(</span><span class="n">x2</span><span class="p">)</span>
    <span class="n">transpose_b</span> <span class="o">=</span> <span class="n">ndim2_orig</span> <span class="o">==</span> <span class="mi">1</span>
    <span class="n">shape_backbone</span> <span class="o">=</span> <span class="n">_check_matmul_shapes</span><span class="p">(</span><span class="n">shape1_orig</span><span class="p">,</span> <span class="n">shape2_orig</span><span class="p">,</span> <span class="s1">&#39;matmul&#39;</span><span class="p">)</span>
    <span class="c1"># infers the shape of the output</span>
    <span class="n">shape_out</span> <span class="o">=</span> <span class="n">shape_backbone</span> <span class="o">+</span> <span class="n">_infer_shape_rem</span><span class="p">(</span><span class="n">shape1_orig</span><span class="p">,</span> <span class="n">shape2_orig</span><span class="p">,</span>
                                                  <span class="n">ndim1_orig</span><span class="p">,</span> <span class="n">ndim2_orig</span><span class="p">,</span> <span class="n">transpose_b</span><span class="p">)</span>

    <span class="n">x1</span> <span class="o">=</span> <span class="n">_expand</span><span class="p">(</span><span class="n">x1</span><span class="p">,</span> <span class="mi">2</span><span class="p">)</span>
    <span class="n">x2</span> <span class="o">=</span> <span class="n">_expand</span><span class="p">(</span><span class="n">x2</span><span class="p">,</span> <span class="mi">2</span><span class="p">)</span>
    <span class="k">if</span> <span class="n">F</span><span class="o">.</span><span class="n">rank</span><span class="p">(</span><span class="n">x2</span><span class="p">)</span> <span class="o">==</span> <span class="mi">2</span><span class="p">:</span>
        <span class="k">if</span> <span class="n">F</span><span class="o">.</span><span class="n">rank</span><span class="p">(</span><span class="n">x1</span><span class="p">)</span> <span class="o">&gt;</span> <span class="mi">2</span><span class="p">:</span>
            <span class="n">x1</span> <span class="o">=</span> <span class="n">F</span><span class="o">.</span><span class="n">reshape</span><span class="p">(</span><span class="n">x1</span><span class="p">,</span> <span class="p">(</span><span class="o">-</span><span class="mi">1</span><span class="p">,</span> <span class="n">shape1_orig</span><span class="p">[</span><span class="o">-</span><span class="mi">1</span><span class="p">]))</span>
        <span class="n">res</span> <span class="o">=</span> <span class="n">P</span><span class="o">.</span><span class="n">MatMul</span><span class="p">(</span><span class="kc">False</span><span class="p">,</span> <span class="n">transpose_b</span><span class="p">)(</span><span class="n">x1</span><span class="p">,</span> <span class="n">x2</span><span class="p">)</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="c1"># broadcasts x1.shape[:-2] with x2.shape[:-2]</span>
        <span class="n">ndim_aligned</span> <span class="o">=</span> <span class="n">_max</span><span class="p">(</span><span class="n">ndim1_orig</span><span class="p">,</span> <span class="n">ndim2_orig</span><span class="p">)</span>
        <span class="n">x1</span> <span class="o">=</span> <span class="n">_expand</span><span class="p">(</span><span class="n">x1</span><span class="p">,</span> <span class="n">ndim_aligned</span><span class="p">)</span>
        <span class="n">x2</span> <span class="o">=</span> <span class="n">_expand</span><span class="p">(</span><span class="n">x2</span><span class="p">,</span> <span class="n">ndim_aligned</span><span class="p">)</span>
        <span class="n">shape1_aligned</span><span class="p">,</span> <span class="n">shape2_aligned</span> <span class="o">=</span> <span class="n">F</span><span class="o">.</span><span class="n">shape</span><span class="p">(</span><span class="n">x1</span><span class="p">),</span> <span class="n">F</span><span class="o">.</span><span class="n">shape</span><span class="p">(</span><span class="n">x2</span><span class="p">)</span>
        <span class="n">x1</span> <span class="o">=</span> <span class="n">_broadcast_to</span><span class="p">(</span><span class="n">x1</span><span class="p">,</span> <span class="n">shape1_aligned</span><span class="p">[:</span><span class="o">-</span><span class="mi">2</span><span class="p">],</span> <span class="n">shape_backbone</span><span class="p">,</span> <span class="n">ndim_aligned</span><span class="p">)</span>
        <span class="n">x2</span> <span class="o">=</span> <span class="n">_broadcast_to</span><span class="p">(</span><span class="n">x2</span><span class="p">,</span> <span class="n">shape2_aligned</span><span class="p">[:</span><span class="o">-</span><span class="mi">2</span><span class="p">],</span> <span class="n">shape_backbone</span><span class="p">,</span> <span class="n">ndim_aligned</span><span class="p">)</span>
        <span class="n">res</span> <span class="o">=</span> <span class="n">P</span><span class="o">.</span><span class="n">BatchMatMul</span><span class="p">(</span><span class="kc">False</span><span class="p">,</span> <span class="n">transpose_b</span><span class="p">)(</span><span class="n">x1</span><span class="p">,</span> <span class="n">x2</span><span class="p">)</span>

    <span class="k">if</span> <span class="n">dtype</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
        <span class="n">res</span> <span class="o">=</span> <span class="n">res</span><span class="o">.</span><span class="n">astype</span><span class="p">(</span><span class="n">dtype</span><span class="p">)</span>
    <span class="k">return</span> <span class="n">F</span><span class="o">.</span><span class="n">reshape</span><span class="p">(</span><span class="n">res</span><span class="p">,</span> <span class="n">shape_out</span><span class="p">)</span></div>


<span class="nd">@constexpr</span>
<span class="k">def</span> <span class="nf">_create_cummin_perm</span><span class="p">(</span><span class="n">axis</span><span class="p">,</span> <span class="n">x_shape</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;Insure axis is in [-len(x_shape),len(s_shape)-1]&quot;&quot;&quot;</span>
    <span class="n">len_axis</span> <span class="o">=</span> <span class="nb">len</span><span class="p">(</span><span class="n">x_shape</span><span class="p">)</span>
    <span class="k">if</span> <span class="ow">not</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">axis</span><span class="p">,</span> <span class="nb">int</span><span class="p">):</span>
        <span class="k">raise</span> <span class="ne">TypeError</span><span class="p">(</span><span class="sa">f</span><span class="s2">&quot;The date type of &#39;axis&#39; should be Int, but got </span><span class="si">{</span><span class="n">axis</span><span class="si">}</span><span class="s2">.&quot;</span><span class="p">)</span>
    <span class="k">if</span> <span class="n">axis</span> <span class="o">&lt;</span> <span class="o">-</span><span class="n">len_axis</span> <span class="ow">or</span> <span class="n">axis</span> <span class="o">&gt;</span> <span class="n">len_axis</span><span class="p">:</span>
        <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="sa">f</span><span class="s2">&quot;The value of axis should be in [</span><span class="si">{</span><span class="o">-</span><span class="n">len_axis</span><span class="si">}</span><span class="s2">, </span><span class="si">{</span><span class="n">len_axis</span><span class="si">}</span><span class="s2">], but got </span><span class="si">{</span><span class="n">axis</span><span class="si">}</span><span class="s2">.&quot;</span><span class="p">)</span>
    <span class="n">prem</span> <span class="o">=</span> <span class="p">[</span><span class="n">i</span> <span class="k">for</span> <span class="n">i</span> <span class="ow">in</span> <span class="nb">range</span><span class="p">(</span><span class="n">len_axis</span><span class="p">)]</span>
    <span class="k">if</span> <span class="n">axis</span> <span class="o">&lt;</span> <span class="mi">0</span><span class="p">:</span>
        <span class="n">axis</span> <span class="o">=</span> <span class="n">axis</span> <span class="o">+</span> <span class="n">len_axis</span>
    <span class="n">prem</span><span class="p">[</span><span class="mi">0</span><span class="p">],</span> <span class="n">prem</span><span class="p">[</span><span class="n">axis</span><span class="p">]</span> <span class="o">=</span> <span class="n">axis</span><span class="p">,</span> <span class="mi">0</span>
    <span class="n">prem</span> <span class="o">=</span> <span class="nb">tuple</span><span class="p">(</span><span class="n">prem</span><span class="p">)</span>
    <span class="k">return</span> <span class="n">prem</span>


<div class="viewcode-block" id="cummin"><a class="viewcode-back" href="../../../../api_python/ops/mindspore.ops.cummin.html#mindspore.ops.cummin">[docs]</a><span class="k">def</span> <span class="nf">cummin</span><span class="p">(</span><span class="n">x</span><span class="p">,</span> <span class="n">axis</span><span class="p">):</span>
    <span class="sa">r</span><span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Computation of the cumulative minimum of elements of &#39;x&#39; in the dimension axis,</span>
<span class="sd">    and the index location of each maximum value found in the dimension &#39;axis&#39;.</span>

<span class="sd">    It returns the cumulative minimum of elements and the index.</span>

<span class="sd">    .. math::</span>
<span class="sd">        \begin{array}{ll} \\</span>
<span class="sd">            y{i} = min(x{1}, x{2}, ... , x{i})</span>
<span class="sd">        \end{array}</span>

<span class="sd">    Args:</span>
<span class="sd">        x (Tensor): The input tensor, rank of `input_x` &gt; 0.</span>
<span class="sd">        axis (Int): The dimension to do the operation, The axis is in the range from -len(`input_x`.shape)</span>
<span class="sd">          to len(`input_x`.shape) - 1. When it&#39;s in the range from 0 to len(`input_x`.shape) - 1, it means starting</span>
<span class="sd">          from the first dimension and counting forwards, When it&#39;s less than 0, it means we&#39;re counting backwards</span>
<span class="sd">          from the last dimension. For example, -1 means the last dimension.</span>

<span class="sd">    Outputs:</span>
<span class="sd">        - **output** (Tensor) - The output tensor of the cumulative minimum of elements.</span>
<span class="sd">        - **indices** (Tensor) - The result tensor of the index of each minimum value been found.</span>

<span class="sd">    Raises:</span>
<span class="sd">        TypeError: If `input_x` is not a Tensor.</span>
<span class="sd">        TypeError: If &#39;axis&#39; is not an int.</span>
<span class="sd">        ValueError:If &#39;axis&#39; is out the range of [-len(`input_x`.shape) to len(`input_x`.shape) - 1]</span>

<span class="sd">    Supported Platforms:</span>
<span class="sd">        ``Ascend``</span>

<span class="sd">    Examples:</span>
<span class="sd">        &gt;&gt;&gt; from mindspore import Tensor, ops</span>
<span class="sd">        &gt;&gt;&gt; import mindspore</span>
<span class="sd">        &gt;&gt;&gt; a = Tensor([-0.2284, -0.6628,  0.0975,  0.2680, -1.3298, -0.4220], mindspore.float32)</span>
<span class="sd">        &gt;&gt;&gt; output = ops.cummin(a, axis=0)</span>
<span class="sd">        &gt;&gt;&gt; print(output[0])</span>
<span class="sd">        [-0.2284 -0.6628 -0.6628 -0.6628 -1.3298 -1.3298]</span>
<span class="sd">        &gt;&gt;&gt; print(output[1])</span>
<span class="sd">        [0 1 1 1 4 4]</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="n">cummin_op</span> <span class="o">=</span> <span class="n">inner</span><span class="o">.</span><span class="n">Cummin</span><span class="p">(</span><span class="n">axis</span><span class="o">=</span><span class="mi">0</span><span class="p">)</span>
    <span class="k">if</span> <span class="n">axis</span> <span class="o">==</span> <span class="mi">0</span><span class="p">:</span>
        <span class="n">out1</span><span class="p">,</span> <span class="n">out2</span> <span class="o">=</span> <span class="n">cummin_op</span><span class="p">(</span><span class="n">x</span><span class="p">)</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="n">transpose</span> <span class="o">=</span> <span class="n">P</span><span class="o">.</span><span class="n">Transpose</span><span class="p">()</span>
        <span class="n">x_shape</span> <span class="o">=</span> <span class="n">P</span><span class="o">.</span><span class="n">Shape</span><span class="p">()(</span><span class="n">x</span><span class="p">)</span>
        <span class="n">prem</span> <span class="o">=</span> <span class="n">_create_cummin_perm</span><span class="p">(</span><span class="n">axis</span><span class="p">,</span> <span class="n">x_shape</span><span class="p">)</span>
        <span class="n">x</span> <span class="o">=</span> <span class="n">transpose</span><span class="p">(</span><span class="n">x</span><span class="p">,</span> <span class="n">prem</span><span class="p">)</span>
        <span class="n">out1</span><span class="p">,</span> <span class="n">out2</span> <span class="o">=</span> <span class="n">cummin_op</span><span class="p">(</span><span class="n">x</span><span class="p">)</span>
        <span class="n">out1</span> <span class="o">=</span> <span class="n">transpose</span><span class="p">(</span><span class="n">out1</span><span class="p">,</span> <span class="n">prem</span><span class="p">)</span>
        <span class="n">out2</span> <span class="o">=</span> <span class="n">transpose</span><span class="p">(</span><span class="n">out2</span><span class="p">,</span> <span class="n">prem</span><span class="p">)</span>
    <span class="k">return</span> <span class="p">[</span><span class="n">out1</span><span class="p">,</span> <span class="n">out2</span><span class="p">]</span></div>


<span class="k">def</span> <span class="nf">resize_nearest_neighbor</span><span class="p">(</span><span class="n">input_x</span><span class="p">,</span> <span class="n">size</span><span class="p">,</span> <span class="n">align_corners</span><span class="o">=</span><span class="kc">False</span><span class="p">):</span>
    <span class="sa">r</span><span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Resizes the input tensor by using the nearest neighbor algorithm.</span>

<span class="sd">    Resizes the input tensor to a given size by using the nearest neighbor algorithm. The nearest</span>
<span class="sd">    neighbor algorithm selects the value of the nearest point and does not consider the</span>
<span class="sd">    values of neighboring points at all, yielding a piecewise-constant interpolant.</span>

<span class="sd">    Args:</span>
<span class="sd">        input_x (Tensor) - The input tensor. The shape of the tensor is :math:`(N, C, H, W)`.</span>
<span class="sd">        size (Union[Tensor, tuple, list]): The target size. The dimension of size must be 2.</span>
<span class="sd">        align_corners (bool): Whether the centers of the 4 corner pixels of the input</span>
<span class="sd">                              and output tensors are aligned. Default: False.</span>

<span class="sd">    Outputs:</span>
<span class="sd">        Tensor, the shape of the output tensor is  :math:`(N, C, NEW\_H, NEW\_W)`.</span>
<span class="sd">        The data type is the same as the `input_x`.</span>

<span class="sd">    Raises:</span>
<span class="sd">        TypeError: If `input_x` is not a Tensor.</span>
<span class="sd">        TypeError: If `size` is neither tuple nor list.</span>
<span class="sd">        TypeError: If `align_corners` is not a bool.</span>
<span class="sd">        ValueError: If length of `size` is not equal to 2.</span>

<span class="sd">    Supported Platforms:</span>
<span class="sd">        ``Ascend`` ``GPU`` ``CPU``</span>

<span class="sd">    Examples:</span>
<span class="sd">        &gt;&gt;&gt; input_tensor = Tensor(np.array([[[[-0.1, 0.3, 3.6], [0.4, 0.5, -3.2]]]]), mindspore.float32)</span>
<span class="sd">        &gt;&gt;&gt; output = ops.ResizeNearestNeighbor(input_tensor, (2, 2))</span>
<span class="sd">        &gt;&gt;&gt; print(output)</span>
<span class="sd">        [[[[-0.1  0.3]</span>
<span class="sd">           [ 0.4  0.5]]]]</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="k">if</span> <span class="n">size</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
        <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="sa">f</span><span class="s1">&#39;For ResizeNearestNeighbor, size could not be None.&#39;</span><span class="p">)</span>
    <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">size</span><span class="p">,</span> <span class="p">(</span><span class="nb">tuple</span><span class="p">,</span> <span class="nb">list</span><span class="p">)):</span>
        <span class="n">resize</span> <span class="o">=</span> <span class="n">P</span><span class="o">.</span><span class="n">ResizeNearestNeighbor</span><span class="p">(</span><span class="n">size</span><span class="p">,</span> <span class="n">align_corners</span><span class="p">)</span>
        <span class="k">return</span> <span class="n">resize</span><span class="p">(</span><span class="n">input_x</span><span class="p">)</span>
    <span class="k">if</span> <span class="n">is_const</span><span class="p">(</span><span class="n">size</span><span class="p">):</span>
        <span class="n">size</span> <span class="o">=</span> <span class="n">size</span><span class="o">.</span><span class="n">asnumpy</span><span class="p">()</span>
        <span class="n">resize</span> <span class="o">=</span> <span class="n">P</span><span class="o">.</span><span class="n">ResizeNearestNeighbor</span><span class="p">(</span><span class="n">size</span><span class="p">,</span> <span class="n">align_corners</span><span class="p">)</span>
        <span class="k">return</span> <span class="n">resize</span><span class="p">(</span><span class="n">input_x</span><span class="p">)</span>

    <span class="n">resize</span> <span class="o">=</span> <span class="n">DynamicResizeNearestNeighbor</span><span class="p">(</span><span class="n">align_corners</span><span class="p">)</span>
    <span class="k">return</span> <span class="n">resize</span><span class="p">(</span><span class="n">input_x</span><span class="p">,</span> <span class="n">size</span><span class="p">)</span>
</pre></div>

           </div>
           
          </div>
          <footer>

  <hr/>

  <div role="contentinfo">
    <p>
        &#169; Copyright 2021, MindSpore.

    </p>
  </div>
    
    
    
    Built with <a href="https://www.sphinx-doc.org/">Sphinx</a> using a
    
    <a href="https://github.com/readthedocs/sphinx_rtd_theme">theme</a>
    
    provided by <a href="https://readthedocs.org">Read the Docs</a>. 

</footer>
        </div>
      </div>

    </section>

  </div>
  

  <script type="text/javascript">
      jQuery(function () {
          SphinxRtdTheme.Navigation.enable(true);
      });
  </script>

  
  
    
   

</body>
</html>